High-dimensional limit theorems for SGD: Effective dynamics and critical scaling
Authors: Gerard Ben Arous
Reza Gheissari
Aukosh Jagannath
Venue: NeurIPS 2022 Accept (2022)
Keywords: stochastic gradient descent
high-dimensional statistics
online algorithms
PCA
XOR
Gaussian mixture model
two-layer network

Abstract: We study the scaling limits of stochastic gradient descent (SGD) with constant step-size in the high-dimensional regime. We prove limit theorems for the trajectories of summary statistics (i.e., finite-dimensional functions) of SGD as the dimension goes to infinity. Our approach allows one to choose the summary statistics that are tracked, the initialization, and the step-size. It yields both ballistic (ODE) and diffusive (SDE) limits, with the limit depending dramatically on the former choices. We find a critical scaling regime for the step-size below which this ``effective dynamics" matches gradient flow for the population loss, but at which, a new correction term appears which changes the phase diagram. About the fixed points of this effective dynamics, the corresponding diffusive limits can be quite complex and even degenerate. 
We demonstrate our approach on popular examples including estimation for spiked matrix and tensor models and classification via two-layer networks for binary and XOR-type Gaussian mixture models. These examples exhibit surprising phenomena including multimodal timescales to convergence as well as convergence to sub-optimal solutions with probability bounded away from zero from random (e.g., Gaussian) initializations. 


High-dimensional limit theorems for SGD:
Effective dynamics and critical scaling
Gérard Ben Arous
Reza Gheissari
Department of Mathematics
Miller Institute
Courant Institute
Departments of EECS and Statisties
New York University
University of California, Berkeley
New York, NY benarous@cims. edu
Berkely, CA gheissari@berkellee edu
Aukosh Jagannath
Department of Statistics and Actuarial Science
Department of Applied Mathematics
University of Waterloo
Waterloo, ON a. jaaaannnthhuuuttrrooo.oooo
Abstract
We study the scaling limits of stochastic gradient descent (SGD) with constant step- size in the high-dimensional regime. We prove limit theorems for the trajectories of summary statistics (i.e., finite-dimensional functions) of SGD the dimension goes to infinity. Our approach allows one to choose the summary statisties that are tracked, the initialization, and the step-size. It yields both ballistic (ODE) and diffusive (SDE) limits, with the limit depending dramatically on the former choices. We find critical scaling regime for the step-size below which this "effective dynamics" " matches gradient flow for the population loss, but at which, a new correction term appears which changes the phase diagram. About the fixed points of this effective dynamics, the corresponding diffusive limits can be quite complex and even degenerate. We demonstrate estimation for spiked matrix and tensor models and classification via two-layer networks for binary and XOR-type Gaussian mixture models. These examples exhibit surprising phenomena including multimodal timescales to convergence as well convergence sub-optimal solutions with probability bounded away from zero from random (e.g., Gaussian) initializations.
approach popular examples including
Introduction
Background. Stochastic gradient descent (SGD) is the go-to method for large-scale optimization problems in data science. It is used train complex parametric models high- high-dimensional data. Since its introduction in (39 there has been a tremendous amount of work in analyzing its evolution.
In fixed dimensions, the asymptotic theory of SGD, and stochastic approximations more broadly, is by now classical. There have been works on path- wise limit theorems, such as functional central limit [6 6 At the core of this line of
theorems and even large deviations principles 39 31 work is the idea that the limit where the size, or learning rate, tends to zero, the trajectory of SGD with fixed loss function (appropriately rescaled in time) converges the solution of gradient flow for the population loss with the same initialization. Recently there has been considerable interest in quantifying the rate of this trajectory- wise convergence to higher order, in terms of diffusion approximation. Namely, there are many works developing asymptotic expansions of the trajectory in
28 22
7
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

the learning rate [29 23 25 II 26. Motivated by this, there is rich line of work bounding the time modifications) There is also an interesting line of work obtaining PDE limits in the "shallow network" regime where the dimension of the parameter space
to equilibrium for the associated diffusion approximation (as well as Langevin-
under uniform ellipticity assumptions 129 36
53
diverges but the dimension of the data remains constant: see e.g.. K2 40 12
55 2
In recent years, there has been considerable interest understanding the high-dimensional setting, where one is constrained in the amount of data or the run-time of the algorithm due to the high- dimensional nature of the data and the complexity of the model being trained. In these regimes, one cannot simply take the learning rate to be arbitrarily small as this would force an unlimited sample size and run-time, This is common issue in nigh-dimensional statistics and the standard analytic approach is to study regimes where the sample size scales with the dimension of the problem [50] (5T.
For SGD with constant learning rate, there has been recent progress quantifying the dimension dependence of the sample complexity for various tasks on general (pseudo or quasi-) convex objectives
211 [51 and special classes of non-convex objectives To] 7 3. There has also been important work on scaling limits as the dimension tends to infinity for the specific problems of linear
regression 2 K55 Online PCA Spp 2 and phase retrieval 47 from random starts, and teacher- student networks 41
4 and two-layer networks for XOR Gaussian mixtures
[77 from warm
starts. We also note that the study of nigh-dimensional regimes of gradient descent and Langevin dynamics have a history from the statistical physics perspective, e.g., in [3 14
30 [0 27
Our contributions. We develop unified approach the scaling limits of SGD in high-dimensions with constant learning rate that allows us to understand a broad range of estimation tasks. One of course cannot develop high-dimensional scaling limit for the full trajectory of SGD the dimension of the underlying parameter space is growing. On the other hand, in practice, one is rarely interested in the full trajectory; instead one typically tracks the trajectory of various summary statistics of the algorithm's 's evolution, such as the loss, the amplitude of various weights, or correlations between the classifier and the ground truth (in a supervised setting). We show in Theorem 2.2 that under mild regularity assumptions, the evolution of these summary statistics converges as the dimension grows to the solution of a system of (possibly stochastic) differential equations. These effective dynamics depend dramatically on the initializations (warm vs. random or cold), the parameter regions in which one is developing the scaling limit, and the scaling of the step-size with the dimension.
In practice, SGD often exhibits two types of phases in training: ballistic phases where the summary statistics macroscopically change in value, and diffusive phases, where they fluctuate microscopically. (During training, the evolution can start with either, and can even alternate multiple times between these phases.) Our approach allows us to develop scaling limits for both types of phases.
In ballistic phases, the effective dynamics are given by an ordinary differential equation (ODE) and intuition that the summary statisties evolve under the gradient flow for the
the imite-dimensionn population loss is correct provided the (constant) learning rate is sufficiently small in the dimension. When the learning rate follows certain critical high-dimensional statistics literature- an additional correction term appears. At this critical scaling, the phase portrait deviates significantly from that of the population gradient flow. Furthermore, in microscopic neighborhoods of the fixed points of this ODE, the effective dynamics become diffusive and are given by SDEs which can exhibit a wide range of (possibly degenerate) behaviors. We note that the appearance of the correction term in the ballistic phase was first observed in the setting of teacher-student networks in 41 4
fiitinngmmtccinngg scalings commonly used in the
and very recently investigated in detail in 40
As a simple, first example of the departure of the effective dynamics in the critical step-size regime from the classical perspective, we study estimation for spiked matrix and tensor models in Section B In these models, the effective dynamics are exactly solvable and when the step-size scales critically with the dimension, in the ballistic phase the dynamics have additional fixed points as compared to the population gradient flow. The stability of these fixed points exhibit sharp transitions at special signal-to- to-so ratios. When initialized randomly, the SGD starts in a microscopic neighborhood of an uninformative such fixed point, within which its effective dynamics become diffusive and exhibit a sharp transition between stable and unstable Ornstein-Uhlenbedd processes.
To demonstrate our approach on more complex classification tasks typically studied using neural networks, we study a Gaussian mixture model analogue of the classical XOR problem in Section 5 (The XOR problem is arguably the canonical example of a decision boundary requiring at least
2

---

two-layers to represent (3..) Here we find that the natural summary statistics are 22 dimensional, and their (ballistic) effective dynamies exhibit a rich phenomenology between some 39 fixed point regions of varying topological dimension. Surprisingly, we find that if we initialize the weights of the network randomly (following a Gaussian distribution), then the algorithm will converge to a classifier with macroscopic generalization error with probability 29/32 and then follow a degenerate diffusion.
Before delving into the XOR problem, we first analyze the classification of two component Gaussian mixture model in Section g regression, but with a two-layer network it exhibits some similar phenomenologies to the XOR problem while being more amenable to finer analysis. Here, we again find that if with random initial weights, with probability 1/2 the SGD will first converge to a classifier with macroscopic generalization error, and then follow a degenerate diffusion in a microscopic neighborhood of that set of unstable fixed points. We demonstrate this both empirically for positive digall-o-nois ratio and theoretically in the limit where the SNR tends to zero after the dimension tends to infinity.
This task is of course best solved using a one-layer network i.e., logistic
2 Main result
Suppose that we are given sequence of i.i.d. data Yi, Y2,... taking values in Vn, C &Rd with law Pa E .M, (, and a loss function La : Xn x Vn, - IR, where here Xn S IRP^ is the parameter space. Consider online stochastic gradient descent with constant learning rate, 6n, which is given by
= 8n VX,((e
with possibly random initialization Xo ~ thn E (Xx)). Our interest is in understanding this evolution, (X²). in the regime where both Pv. and d, + (i as n x. To this end, suppose that there is finite collection of summary statistics of (Xt) whose evolution we are interested in. More precisely, suppose that we are given a sequence of functions u, E C (RPr; RR*) for some fixed k, where un (2: = (uf (x)...., u2 (z)), and our goal is to understand the evolution of u,,
To develop scaling limit, we need some assumptions
the relationship between how the step-size scales in relation to the loss, its gradients, and the data distribution. To this end let H (x, Y) = Ln (zr, Y)- ((x), where ((x) = L, (zr, Y)J. Throughout the following, we suppress the dependence of H on Y and simply write H (x)., and instead view H as random function of æ. Definition 2.1. We say that a triple (u,, In,, Px) is 6, -o-aalaliz if there is an exhaustion by compact sets of R&*, call it (Ex) x, and constants 0 < C(K) < oo (independent of n) such that
V' ullop S C(K) on 5-3-5)/2 for j = 2.3;
1. max SuPweu, ' (Ex)
2. SUPree, (Ex) | < C(K), and
3. maxigiçk Sup (Ex)
E (V H. Vu,) S (K)),,2 and
E ( u, VH & VH - vy = 1.
maxişish sup sup.eu (Ex)
When these hold we call the sequence ( (Fx) the localizing sequence of (un, Ie;; Pn).
Localizability is wider than uniform Lipchitz or smoothness assumptions common to the literature. In particular, it does not imply that the population loss is Lipschitz everywhere. as we may have that UK u, (E,)(E)) does not cover IR",, nor does it imply uniform smoothness of L as we will be taking 8n 0 with n. To motivate the scaling relations between 2-3, note that if 6, x p corresponding to linear sample complexity, the scaling relation is the same as what one would get e.g., if H were a random vector with independent entries of bounded variance in R and e.g., Vu; = 5, 1/2
We now turn to the statement of our main result. Let Pa denote the space of positive semi-definite k k matrices and for a function and measure μ we let f.j denote the push-forward of μ. Let Jn = (Vu)) denote the Jacobian of the summary statistics. Also, let V(x) = EE (x) 8 (=)) denote the covariance matrix for Hl a point x and define the corresponding first and second-order differential operators.
An = - 0,00;
and
Cn
1 - VijO,Og.
We then have the following convergence result.
3

---

Theorem 2.2. Let xx be the SGD initialized Xo ~ for pn E (R with learning rate on for the loss -.l.,-) and data distribution P.. Suppose that u, is such that the triple
(un, T.u Pa) is on.loocalizbble with localizing sequence (Ex). Suppose furthermore that there exists locally Lipschitz h R + I
and 2 R* + Pki such that for every K,
sup ll Sn un (x) (((u(()))] +0
(2.1)
sup J6 - I(u, (2))1l -
(2.2)
where the suprema are over x: E (Ex). Then if we let (un(t)) be the linear interpolation of (u, (X 6,n [tó,. 1)) and the initial data Hos is such that (un)+ + v weakty, then, (u,,t)); + (uz)
weakly as n - oc, where u, is the solution to
du, = h(u + dB.
(2.3)
initialized from v, where B, is a standard Brownian motion in R
Theorem 2.2 is proved in Appendix C
The proof can be seen version of the martingale problem for high- digh-dimmnsional SGD.
(see 46
2.1
Comparison to fixed dimensional perspective: critical v.s. subcritical step-sizes
Let us compare this with the classical limit theory of SGD in fixed dimension. For the sake of this discussion, suppose that not only does (2.1) hold, but each of the two terms A. u and 8, u individually admit n + limits: namely that there exists f,g IR* + IR* such that
sup [lAA,u()) f(u, (x))I v ||,, Cn (x) (x))l - 0, zeu. (Exx
(2.4)
in which case, evidently 2.1 holds with h f + g. When 2.4 both hold, we call f, gg and the population drift, the population corrector, and the diffusion matrix of u respectively.
From the fixed dimensional perspective, when 2.4 holds, one predicts to asymptotically solve
du,- = d(,
(2.5)
with initial data Ug ~ u.ft as this is its evolution under gradient descent on the population loss क. Evidently this perspective only applies in the high-dimensional limit of Theorem 2.2 if both the population corrector g and the diffusion matrix E are zero. We find that for any triple (u,, Ln. Pn). there is scaling of the learning rate dn with n below which g = Σ = O, and the effective dynamics agree with the population dynamics (2.5)call call this the sub-critical scaling regime, where the classical perspective applies-and a critical scaling regime in which g and Σ may be non-zero, and the dimensionnaltty induces non-tn tonvirivial corrections to f. (In the case of teacher-student networks, the terms f and g can be compared to the "learning" and "variance" terms in Eq. (14a) of 49 .)
To see this, notice that if the triple (u,,, Ln, P.) is oocalaccleeaele for some 6, - O, then it is also 8, localacaizeelle for every sequence 6', = ((,)). J.. If furthermore 2-4)- (.2) hold for 8, with some f, f,g and Σ, then these limits also exists for 8, = 0(5,)) with the same but with g Σ 0. As such, there can be exactly one scaling of 5n with n at which or Σ may be non-zero, and for all smaller scales of learning rate, the fixed-dimensional perspective of (.5) applies
2.2 Ballistic vs. difflusive behavior of effective dynamics
In all of our examples, the diffusion matrix for the effective dynamics of the most natural choice of summary statistics is zero even in the critical scaling regime where o f. We call this the ballistic limit. In this case, the effective dynamics of the summary statistics is given by the ODE system
du,= = h(,,ddt.
(2.6)
In these settings, the phase portrait of the summary statistics is asymptotically that of this flow.
By construction of the scaling limit, the phase portrait of the ballistic limit only describes the evolution of summary statistics on length-scales that are order and time-scales that are order 1/6,.. If one is
Note that if ón o(6%), then limiting g, may not exist for 8',, there is no super-critical regime.

---

then interested in the evolution of un in microscopic o(1) neighborhoods of the fixed points of 2.6 Theorem 2.2 also allows one develop separate diffusive limits there.
To understand diffusive regimes, one must apply Theorem 2.2 to a re-centered and re-scaled version of the summary parameters, (t) 5,-² (u., (t) u.)) where u, is fixed point of (2.6) To apply
Theorem 2.2, a must be chosen appropriately so that the triple (t), Ln, Pn) is o₂ localizable and to pick out the next order drifts for ü-the first order term being zero microscopically close to u, und such that the initial data still converges (ūn) (un).nn + V.
This then leads to the rescaled effective dynamics of the summary statistics u, near u,
dū, = h(ū)ddt + E (1,)(,,
with un ~ V.
(2.7)
The rescaled effective dynamics are similar in spirit to diffusion approximations typically one finds for the evolution of SGD near critical points in fixed dimensions. However, we note two important differences as compared to this perspective. Firstly, since this is a high-dimensional limit of general summary statistics, (2.7 applies in a neighborhood of a fixed point of the effective ODE system (2.6) rather than the population dynamics (.5 Secondly, in many examples (indeed all the ones we study) the SDE's s we get are degenerate to some degree, so that uniform ellipticity assumptions typically used to understand hitting and mixing times in these regimes do not apply.
3 Matrix and Tensor PCA
As our first example, we consider the problems of spiked matrix models and spiked tensor models 38 using SGD. These examples are exactly solvable and only require two summary statistics, correlation observable and a radial term. Even with this relative simplicity, we encounter a wide range of ODE and SDE limits. Interestingly, by means of these SDE limits, we can sharply identify the signal-to-noise noise thresholds for solving the recovery problem by means of the SGD.
Suppose that we are given data of the form Y = ak + W where W is an i.i.d. Gaussian k-tensor, R² is unit vector, and X An 0 is the signal-to-noise ratio. Our goal is infer v. We take as loss the (negative) log-likelihood namely, L(x, Y) = | - 122 The pair of summary statisties m m(x)= = (xr, and r2 r1 (z) = | mu|
m² are such that (xx) 22m + (r + m + c, and the law of L only depends on them: see Section D.1
For the pair u, = (u1, u2) = (m,r1). Theorem 2.2 yields the following effective dynamics. In our normalization with À > 0 fixed, the regime 6n = o(I/n) is sub-critical and the regime 8, = (1/nn) is critical; we focus this normalization for presentation, but note that with different scalings of Xn. the critical learning rates change. For notational simplicity, let R² == m² +1².
Proposition 3.1. Fix k 2 2, X> 0, > 0 and let 8, = -6//1 Then un (l) converges the solution of the following ODE initialized from litmn- un,) )on:
to
ii 2u (Akut k-2 12222-2))
(3.1)
2 = - -KRR (u2 - cs)
We are able to identify and classify the set of fixed points of this effective dynamics. (Recall that the dynamics transits ballistically between these fixed points in 6%_² many steps.) We focus the critical step-size regime with G = 1 where one sees from (3.1) that U2 - 1, which is where a random vector in Rn lies, and where the problem in the matrix case is most directly related to eigenvalue problem (see Appendix D for the generic co dependencies).
Proposition 3.2. Eq. (..) has isolated fixed points classified as follows. Let (r. be as in D.4 and my (k, X) < m. (ki, 2) be as in (.5) (if l = 2, X = and m, = mm = Vd I):
1. An unstable fixed point at (0,0) and a fixed point at (0, if = 2, (0, 1) stable if
X < Xx(2) and unstable if X > (2); if k > 2 (0, 1) is always stable.
"One might also wish rescale time like 8-8, where 3 may depend on t; we leave this future work. 3 Note that one might also add additional penalty terms. The case a ridge penalty is treated in Section 'The sub-critical regime of on o()/)) recovered by sending cs 0 in the below.
5

---

210
115
140
090
(b)
(c)
(d)
Figure 1: Matrix PCA summary statistics in dim. n 1500 run for 10n steps at À = 0.8 < Ae in (a)-(b) and À 1.2 > A₂ in (c). Here, X and mark the stable fixed points of the systems. (a) and (c) demonstrate the stable and unstable OU processes that arise diffusive limits of the m variable, and (b) and (d) depict the trajectories in (m, space.
2. 1f >> Ak(k)): when k = 2, two stable fixed points at (=mm(2), 1). When k 2 3, two unstable fixed points (+:my(k) 1) and two stable fixed points at (#mm, (k:). 1).
Remark 1. The presence of two pairs of fixed points when k 2 3 with non-zero correlation with may seem surprising-ine inddedd indicates that even some warm starts will fail to attain good correlation with the signal when À is finite. This is an interesting consequence of the corrector in (3.1) and if one tracked the es dependence in the above, the fixed point my goes to zero as cs + and this barrier to recovery from warm starts vanishes as one approaches sub-critical step-sizes.
-nndeed
Let us now consider a rescaling of u, in a microscopic neighborhood of the saddle set m = 0. This captures the initial phase from a random start: if jr, NO, /n/n), then (n + ō(o,1) weakly, Now rescale and let un (Vrm, ri). Evidently,
lim,, (a -O 1) @ ö1. Proposition 3.3. Fix k > 2, X >>0 0 and 8n 1/n. Then ú, converges as to the solution of the following SDE initialized from D:
dug 2û, (2A1 kūzz +
(3.2)
dūg 4kūb (ū2 1)dt.
We that ü2 solves an autonomous ODE which converges exponentially 1. When k: = 2, the equation for ún then converges to 4(A 1) dt + 2/2dB, for large t. This is OU process which stable when À < 1 and unstable when A 1. By stitching together the prelimits of these OU processes at sequence of scales interpolating between that of ü, and un one could in principle establish that for any À > 1, SGD reaches the stable fixed points at (=m,(2), (2). 1) in (n log n) steps (with precise asymptotics, etc.), while when À the mean-reverting nature of the OU suggests needs a much larger number of samples in order to correlate with the vector v. See Figure for numerical verification of this intuition. When i: 3, the tensor PCA problem is known be hard for SGD to solve without a polynomially diverging sample complexity or À [3] Accordingly, when is kept finite in n, the expression for ū in (32 is always stable OU-type process. Interestingly, one can also capture the (diverging) signal-to-noise threshold for SGD to recover o in tensor PCA by our methods. Indeed, for 2 3 if one considers A (k-2)) (matching the predicted gradient- algorithm threshold from 4 ) ū, would instead converge to the solution of
dū 2ū, (kA - )tt + ((
dū2 - 4k- k-1 - 1)dt,
which transitions between stable and unstable OU processes at (k) 1, in the matrix case.
4 Two-layer networks for classifying a binary Gaussian mixture
As our second example, we consider the problem of supervised classification of a binary Gaussian mixture model (binary GMM) using two-layer network. Our goal here is to demonstrate how our approach can be used to analyze the performance of SGD for multi-layer networks, and indeed we will find the calculations here be relevant in Section 5 where we consider XOR-type GMM's.
Let us now formalize the problem. Suppose that we are given i.i.d. samples of the form Y (y, X), where y is a Ber(1/2) random variables and, conditionally y. we have X NN(22 L)μ, 1/A).

---

where P E IN is fixed unit vector, is the identity on R and À > 0 is the Segmal-to-niin ratio. Here, y the class label and is the data.
For the sake of concreteness, we consider classification via the following architecture (though our techniques generalize to other settings mutatis mutandis): The first layer has weights (Wi, W2) R x I and ReLu activation, (xx) = I V O; and the second layer has weights v1, E R and sigmoid activation, (z) 1/(1+e r) Our parameter space is then Xx IR and we therefore
take = 2N + 2 when applying Theorem 2.2 As we are interested in supervised classification, we take the usual binary cross- pross-entropy loss with 8 regularization,
L( Wi) ()) = yo g(WX) + log(1 + +e e + p(e, W),
(4.1)
where is applied component wise and pvu W) = (a/2) + +||]))
It can be shown (see Lemma that the law of the loss given point, (u, W) E 2r depends only on the summary statistics,
Up (U1, U2 mı, m2, Rn Riz R22),
(4.2)
where m, = We j and R. = W W with W, = W m,ju denoting the part of W; orthogonal to j.. For point, (v, W) E Nn let
A = E[X x.x.xllwwe x
f)],
A E[X W, 1 1ww.x2ollle - -)],
(4.3)
Bi E E[lw, x22lo
y)
By Lemma E.1
these are functions only of un. and we denote them as such, e.g., A A (u,).
The critical scaling for 5 is then of order (11/n) and we obtain the following effective dynamics.
Proposition 4.1. Let un be as in and fix any X >> 0 and 5, = cs/N. Then un,() (t) converges to the solution of the ODE system,
from f(u) + g(u)). initialized from lim (u,), Pro with:
f m,A," (u) + (u) a0, f.: v:A[(u) (u) + am,,
JR. = A (u) +
+ Riy
and correctors Iv. 9m. O, I
c8 Bij for i.j = 1,2.
Due to the Gaussian integrals defining f. f.g. it is difficult to analyze the ODE system defined by Proposition 4.1 let alone any rescaled effective dynamics. For ease of analysis, we next send À oo small noise regime for the Gaussian mixture. We emphasize that this limit is taken after n >> no and therefore still approximately on the critical scale of X = +(1) which there is a transition in the existence of any fixed point which is a good classifier. In particular, if X An is any diverging sequence, then the limiting effective dynamics would exactly match that attained by now sending À + oc. In Figure 2 we demonstrate numerically that the following predicted fixed points from the + oo limit match those arising at finite large n and À
corresponding
Proposition 4.2. The X oc limit of the ODE system of Proposition is given by
w: 0 (v- m) am:;
m > 0
2
2 0 (-oams) - omT4 else
a (-v m) c09
m1m2
0:
( v.m;)) 2
00;;
else
and R, -2. The fixed points of this system classified as follows. All fixed points have
R = and m, = 0, for i,j = {1,2]. In (v1, 12), the coordinates are classified by
1. A fixed point at (vi, 02) = (0, 0) that stable if > 1/4.
For large X. this is indeed quantitative approximation as f, exhibit locally Lipschitz dependence on -1 the corresponding dynamics converges
by classical well- oolldossddeess results (see, e.g.. 48
7

---

34
(a)
(b)
(c)
(d)
Figure 2: Binary GMM in dim. N = 250 with À 100 and 0.1. (a) fraction of runs converging
to the optimal classifier (x marking the predicted //2 fraction), and (b) endpoints of (m1, m2) in 200 runs (x denoting the À predicted stable fixed points and the unstable rings). (c)-(d) diffusive limits, first for mi, and then for the pair m,, m2 where the diffusion can be seen to be of rank 1.
2. If < 1/4, two unstable sets of fixed points at the quarter-circles given by (VI, v2) having V/4/ > 0 such that vf + v2 = Co for Co == -loggl1 2(x) log(2()).
3. If < 14,, two stable fixed points at (v1, 02) equals ( (VCa, - -VCa and (-VCas VC.)
If Hm, is e.g., given by (2), (0, 12) and W1, W2 N IN (t/((N)) then v == lim((,)) ) Un is (0, In) in the U1 02 coordinates, and is in the basin of attraction of the quarter-circles of item (2) with probability 1/2 and the basin of attraction of the stable fixed points of (3) with probability 1/2.
Let us pause to interpret this result. The stable fixed points when a < 1/4 are the optimal classifiers, whereas the unstable set of fixed points given by item (2) misclassify half of the data. Therefore, the above indicates that when solving the above task with randomly initialized weights, one of the following two scenarios occur, each with probability 1/2 (w.r.t the initialization): the algorithm will converge to the optimal classifier in linear time or will appear to have converged to macroscopically sub-optimal classifier the same timescale, see Figure (b)).
then natural ask about the behaviour of the SGD the latter regime, after it converges to the sub-optimal classifiers which lie on the aforementioned quarter-circles. Proposition 42 rigorously justified the exchange of n - 0o and - oo limits in the ballistic phase. the diffusive phase, one could in principle find the quarter circle of fixed points of the ODE in Proposition 4.1 and consider rescaled observables Ty m, corresponding to blowing up U, m. in diffusive O(n 1/2))/2)))0) neighborhoods about them to get SDE limits from Theorem 2.2 In order to have explicit formulae, in what follows, we consider the diffusive limits obtained when taking X = oc, for which we know the precise locations of these fixed points from Proposition This also captures the limit obtained by taking any An diverging faster than Onn2/2); the numerics of Figure (d)) demonstrate its qualitative consistency with the behavior in microscopic neighborhoods of fixed points even at À finite.
Proposition 4.3. Let 8n = I/N, (a1, E IR, be such that of + Ca and let é; = VN(e, - -a)) and m, VNN a.). When X oc, the SDE system obtained by applying Theorem 2.2 ú, is
dū; = am ū))) + a;(a 2a²) Σ ax + mx)) - 60g
dm, = alē; m.) + a:(a 2a²)
Σ (xx + mm) + 1/2 dB. Emp
R. -
R 22x 2
where E is a constant matrix whose only non-zero entries are En,, Enjryy = Einm = a a²a;aa
Notice that this diffusion matrix is rank 1, so this diffusion is non-trivial but degenerate even in the rescaled coordinates (ū;, m.). Moreover, the entries of E vanish on the axes a 0 or d2 = 0. In particular, crossing from the unstable quarter ring into the quadrants U1 02 < 0 where the stable fixed points lie is impossible in the noiseless setting, and happens much larger timescale at finite A.
5 Two-layer networks for the XOR Gaussian mixture
We end by applying our methods to supervised learning for an XOR-type Gaussian mixture model in g with two-layer network. The ballistic phase of this problem was recently studied in 337
8

---

Me
200
100
309
140
240
340 040 140
(a)
(b)
(c)
(d)
Figure 3: XOR GMM in dim. N 250 with À 1000 and a 0.1. (a) Fraction of runs converging
to the optimal classifier (x marking the predicted 29/22 and 8/52) and (b) endpoints of (mf m² in 200 runs (x denoting the oc predicted stable fixed points, the unstable rings). (c) display the rank-2 diffusive limits in the regime of Proposition 5.1 in (m mm)) and (mf mÿ) coordinates.
Suppose that we are given i..d. samples of the form Y = X)., where is Ber(1/2) as before and X has the following distribution: if y then X is 1/2-1/2 mixture of N(,, and (- I/)) and if = 0 it is 1/2-1/2 mixture of (M, T/A) and .V(v,, T/A), where À > 0, and μl, are orthogonal unit vectors. Here, y the class membership label and X is the data.
Consider the corresponding classification problem using a two-layer neural network, x(X) = (v g( where a and g are as in Section | We take W to be a N matrix and v to be 4-vector, and consider the binary cross-entropy loss as before. It is shown in Lemma F. that the law of the loss point (v, W' depends only on the following 22 variables: for 1 S S 5
Us+.
mf = W We-..
m = W. V
R = W.- W
(5.1)
where W.- = W m Jr m v is the part perpendicular Jr, v. With the choice of u, given by these variables, for any fixed X > 0, the localizability criterion of Definition Pefinitionn..Jaa [an be verified to hold as long as 6, (1/nn). In particular, we can apply Theorem 2.2 to obtain limits in both the ballistic and diffuuive phases. For the precise equations in the ballistic phase, see Proposition
The fixed points of the ballistic dynamics, again in the limit À + co after + oo, are classified as follows (see Proposition F.3) If a > l/8, then the only fixed point is at u, = 0. If < a < 1/8, then let (Io. I In I I. be any disjoint (possibly empty) subsets whose union is {1..... 4). Each such partition corresponds to a connected component of fixed points. Corresponding to such tuple, the connected component of fixed points has R2 for all j, and
1. mf m 82 0 for i lo.
2. m = 0; > such that Eier, os = - logit( (4a) and my = 0 for all i E Tt.
3. m th 0. > 0 such that Eiel, v logit (4a) and m. 0 for all e 4,-
4. m U1 such that Eier, uf - logit (a) and m 0 for all i e It.
5. m = 0/ < 0 such that Eier. of = logit 4cx) and mn = 0 for E 75.
There are 39 connected components of fixed points, 4! of which are stable, one for each permutation where 1t+ 1,,,,,1 I,- are all singletons. For the proof of this limit, see Proposition E3
We can also compute the probability that the effective dynamics in the ballistic phase converges to a stable fixed point (as opposed to an unstable one). From Gaussian initialization fh, where 82 ~ (O 1) and W. ~ N/00, IN /N) independently. this probability will converge to 3/32.
As an example of the diffusions that can arise in the rescaled effective dynamics at the unstable fixed points, let us consider the unstable fixed points in which v has the correct signature (two positive, two negative) but for each of those we are at corresponding quarter-ring. Here, the dynamics effectively becomes a pair of 2 two-layer GMM's on quarter-rings (as in J. that are anti anti-correlated. More precisely, let (a1,u: a2,..) be such that af + a² C. and (a3,r, as,z) such that az + az,, Ca. for Ca logit((4x) Take as fixed points about which we expand to be 0 = m. = ai > 0 and = m. = aip < 0 for i = 3, 4. Namely, we let
V N (w: ai.p))
i= 1,2
J =
VN(u: - ai
i= i=3,4

---

and let
m v VN(m 0i..) i= 1,2
m = VN(m" - di.z) i= 3,4.
(We set m = 0 for i = 1, and m = O for i == 3,4 in û, effectively removing those variables.)
Proposition 5.1. Let 8, = I/N and let ü, = (ü; m my. R. When À = oc, Theorem 2.2 can be applied and (1) converges to the solution of the SDE dū(l) = li(ū)dl + / /((ū)BB where
hag (a(ū: m( ai 4a²) Σ Ek=1, RR. (Ūx + mi))
1,2 (d m - aipl(a - 4a²) Ek=s, aR. + mmx) = 3, 4
bag (resp., hag)) is like ha, for i = 1,2 (resp., i = 3, 4) with ūn and m (resp., m swapped, hRz = 2... and E the constant rank-2 matrix whose non-zero entries are
Eo,e, Enpiing = Eormg = 3a ai
if E (1,2),
Di,e, = Eipg = Eiiny = 3a 2 an
ifi,j E (3,4),
Eo,e, Empma Empres Eerny = -a- -a'aa ai,
i (1.2),} E (3,4).
Acknowledgements
The authors thank the anonymous referees for their careful readings of the manuscript and helpful comments. R.G. thanks the Miller Institute for Basic Research in Science for its support. A... acknowledges the support of the Natural Sciences and Engineering Research Council of Canada (NSERC). Cette recherche a été financée par le Conseil de recherches en sciences naturelles et en génie du Canada (CRSNG), [RGPIN-2020-045] DGFCR-20000001999,
References
[1] Andreas Anastasiou, Krishnakumar Balasubramanian, and Murat A Erdogdu. Normal ap- proximation for stochastic gradient descent via non-asymptotic rates of martingale CLT. In Conference on Learning Theory, pages 115-137. PMLR, 2019.
[2] Dyego Araújo, Roberto Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint arXiv:1906.001 2019.
[3] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from dighesiinnsi inference. Journal of Machine Learning Research, 22(106):1-5,, 2021.
[4] Gérard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Algorithmic thresholds for tensor PCA. Annals of Probability, 48(4):2052-200 2087, 2020.
[5] Gérard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Bounding flows for spherical spin glass dynamics. Communications in Mathematical Physics, 17113):0011-1048 2020.
[6] Michel Benaïm. Dynamics of stochastic approximation algorithms. In Séminaire de Probabilités, XXXIII, volume 1709 of Lecture Notes in Math., pages 1-68. Springer. Berlin, 1999.
[7] Albert Benveniste, Michel Métivier, and Pierre Priouret. Adaptive algorithms and stochastic appreximations, volume 22 of Applications of Mathematics (New York). Springer-Verlas Verlag, Berlin, 1990. Translated from the French by Stephen S. Wilson.
[8] Léon Bottou. On-Line Learning and Stochastic Approximations. Cambridge University Press, USA, 1999.
[9] Léon Bottou and Yan Le Cun. Large scale online learning. In S. Thrun, L. K. Saul, and B. Schölkopf, editors, Advances Neural Information Processing Systems 16, pages 217-224. MIT Press, 2004.
10

---

[10] Michael Celentano, Chen Cheng, and Andrea Montanari. The high- sish-ddimnsionall asymptotics of first order methods with random data. arXiv preprint arXiv:2112.07572. 2021.
[11] Xiang Cheng, Dong Yin, Peter Bartlett, and Michael Jordan. Stochastic gradient and Langevin processes. In International Conference on Machine Learning, pages 1810-1819. PMLR, 2020.
[12] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.
[13] A Crisanti, H Horner, and H-J Sommers. The spherical p-spin interaction spin-glass model. Zeitschrift für Physik Condensed Matter, 92(2):257-217 1993.
[14] Leticia F. Cugliandolo and Jorge Kurchan, Analytical solution of the off-equilibrium dynamics of long-range spin-glass model. Phys. Rev. Lett., 71:173-176, Jul 1993.
[15] Aymeric Dieuleveut, Alain Durmus, and Francis Bach. Bridging the gap between constant step size stochastic gradient descent and Markov chains. Ann. Statist., 1000)1348-8388 06 2020.
[16] Marie Duflo Algorithmes stochastiques, volume 23 of Mathématiques & Applications (Berlin) |Mathematics & Applications). Springer- Verlag, Berlin, 1996.
[17] Paul Dupuis and Harold J Kushner. Stochastic approximation and large deviations: Upper bounds and w.p.1 convergence. SIAM Journal on Control and Optimization, 1989.
[18] Stewart N. Ethier and Thomas G. Kurtz. Markov processes. Wiley Series in Probability and Mathematical Statistics: Probability and Mathematical Statistics. John Wiley &amp: Sons, Inc., New York, 1986. Characterization and convergence.
[19] Rong Ge, Furong Huang, Chi Jin, and Yang Yuan. Escaping from saddle points - online stochastic gradient for tensor decomposition. In Proceedings of The 28th Conference on Learning Theory, volume 40 of Proceedings of Machine Learning Research, pages 797-842, Paris, France, 03-06 Jul 2015. PMLR.
[20] Sebastian Goldt, Madhu Advani, Andrew Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019.
[21] Nicholas J. A. Harvey, Christopher Liaw, Yaniv Plan, and Sikander Randhawa. Tight analyses for non-smooth stochastic gradient descent. In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Thirty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pages 1579-1613, Phoenix, USA, 25-28 Jun 2019. PMLR.
[22] Harold Kushner. Asymptotic behavior of stochastic approximation and large deviations. IEEE transactions automatic control, 29(1 1):984-9900 990, 1984.
[23] Chris Junchi Li, Mengdi Wang, Han Liu, and Tong Zhang. Diffusion approximations for online principal component estimation and global convergence. Advances in Neural Information Processing Systems, 30, 2017.
[24] Chris Junchi Li, Zhaoran Wang, and Han Liu. Online ICA: Understanding global dynamics of nonconvex optimization via diffusion processes. In D. D. Lee, M. Sugiyama, U. Luxburg. I. Guyon, and R. Garnett, editors. Advances in Neural Information Processing Systems 29, pages 4967-4975. Curran Associates, Inc., 2016.
[25] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. The Journal of Machine Learning Research, 20( 14244115-5,, 2019.
[26] Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (SDEs). Advances in Neural Information Processing Systems, 34, 2021.
11

---

[27] Tengyuan Liang. Subhabrata Sen, and Pragya Sur. High-dimensional asymptotics of Langevin dynamics in spiked matrix models. arXiv preprint arXiv:2204.04476, 2022.
[28] Lennart Ljung. Analysis of recursive stochastic algorithms. IEEE Trans. Automatic Control, AC-22(4):551-575, 1977.
[29] Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as approximate Bayesian inference. The Journal of Machine Learning Research, 18(1):487344907, 2017.
[30] Stefano Sarao Mannelli, Giulio Biroli, Chiara Cammarota, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Marvels and pitfalls of the Langevin algorithm in noisy high- dimensional inference. Physical Review X, 10(1):01 1057, 2020.
[31] D. L.. McLeish. Functional and random central limit theorems for the Robbins- process. Journal of Applied Probability, 13(1), 1976.
[32] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proc. Nat.. Acad. Sci. USA, 115(3)::8666--7677) 2018.
[33] Marvin Minsky and Seymour Papert. An introduction to computational geometry. Cambridge tiass., HIT, 479:480, 1969.
[34] Deanna Needell, Nathan Srebro, and Rachel Ward. Stochastic gradient descent, weighted sampling, and the randomized kaczmarz algorithm. In Proceedings of the 27th International Conference on Neural Information Processing Systems- Volume 1, Cambridge, MA. USA, 2014.
MIT Press.
[35] Courtney Paquette, Kiwon Lee, Fabian Pedregosa, and Elliot Paquette. SGD in the large: Average-case analysis. asymptotics, and stepsize criticality. In Conference on Learning Theory, pages 3548-3626. PMLR, 2021.
[36] Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex learning via stochastic gradient Langevin dynamics: nonasymptotic analysis. volume 65 of Proceedings of Machine Learning Research, pages 1674-1703, Amsterdam, Netherlands, 07-10 Jul 2017. PMLR.
[37] Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová. Classifying high- dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference on Machine Learning, pages 8936-8947. PMLR, 2021.
[38] Emile Richard and Andrea Montanari. A statistical model for tensor PCA. In Advances in Neural Information Processing Systems, pages 2897-2905, 2014.
[39] Herbert Robbins and Sutton Monro. A stochastic approximation method. Ann. Math. Statistics, 22:400-407, 1951.
[40] Grant M Rotskoff and Eric Vanden-Bijnden. Trainability and accuracy of neural networks: An interacting particle system approach. arXiv preprint 1rXXv:1805.0091: 2018.
[41] David Saad and Sara Solla. Dynamics of on-line gradient descent learning for multilayer neural networks. Advances in neural information processing systems, 8, 1995.
[42] David Saad and Sara A Solla. On-line learning in soft committee machines. Physical Review E, 52(4):4225, 1995.
[43] Stefano Sarao Mannelli, Giulio Biroli, Chiara Cammarota, Florent Krzakala, and Lenka Zde- borová. Who is afraid of big bad minima? analysis of gradient-flow in spiked matrix- -tesor models. Advances in Neural Information Processing Systems, 32, 2019.
[44] Ohad Shamir. Convergence of stochastic gradient descent for PCA. In International Conference on Machine Learning, pages 257-265, PMLR, 2016.
[45] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: central limit theorem. Stochastic Processes and their Applications, 1300 130(3):1800-1852 2020.
12

---

[46] Daniel W. Stroock and S. R. Srinivasa Varadhan. Multidimensionall diffussion processes. Classics Mathematics. Springer- Verlag, Berlin, 2006. Reprint of the 1997 edition.
[47] Yan Shuo Tan and Roman Vershynin. Phase retrieval via randomized Kaczmarz: theoretical guarantees. Information and Inference: A Journal of the IMA, 8(1):97-123, 04 2018.
[48] Gerald Teschl. Ordinary differential equations and dynamical systems, volume 140. American Mathematical Soc., 2012.
[49] Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. arXiv preprint arXiv:2202.0029)) 2022.
[50] Roman Vershynin. Digge--Dimssinall Probability. Cambridge University Press (to appear),
2018.
[51] Martin Wainwright. Iligh-dimensional statistics; non-asymptotic viewpoint, volume 48. Cambridge University Press, 2019.
[52] Chuang Wang, Jonathan Mattingly, and Yue Lu. Scaling limit: Exact and tractable analysis of online learning algorithms with applications to regularized regression and PCA. arXiv preprint 1rXiv:1712.04332. 2017.
[53] Yuchen Zhang, Percy Liang, and Moses Charikar. A hitting time analysis of stochastic gradient Langevin dynamics. volume 65 of Proceedings of Machine Learning Research, pages 1980- 2022, Amsterdam, Netherlands, 07-10 Jul 2017. PMLR.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] See appendix
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] See supplemen- tal material
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [No] Experiments are just illustrative single runs
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [No]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any assets either in the supplemental material as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "vVVxM2OR6HZ", "reviewer": "Reviewer_hQ5R", "summary": "This work investigates the high-dimensional limit for the evolution of summary statistics under one-pass stochastic gradient descent (SGD).\n\nGiven some data, a loss function and a learning rate, the one-pass SGD dynamics is defined by evaluating the gradient *at a single data sample* and performing a gradient step in this direction. In the so-called *classical* limit where the dimensions of the problem are fixed and the learning rate is taken to zero, one-pass SGD is known to converge to the gradient flow in the population risk. This work characterises the corrections to the gradient flow dynamics arising when the dimensions of the problem (data dimension, number of parameters) scale relatively to the learning rate.\n\nIts main theoretical contribution is to map the one-pass SGD dynamics into a set of tractable low-dimensional stochastic differential equations (SDE) for the evolution of the summary statistics (i.e. low-dimensional functions of interest of the parameters, such as their covariance and correlation with a ground truth).\n\nExplicit examples such as planted matrix and tensor PCA and Gaussian mixture classification are discussed in detail, where it is shown how the SDE simplifies in different regimes associated to different choices of scalings for the dimensions.", "questions": "-**[Q1]**: In L137, \"fixed dimension\" means both fixed $d_{n}$ and $p_{n}$?\n\n-**[Q2]**: In [49] it was shown that for teacher-student two-layer NNs by letting the hidden-layer width grow faster than the data dimension, one can go from a regime of imperfect to perfect learning (corresponding to the cross-over between a critical to sub-critical regime in the author's vocabulary). Can a similar conclusion be drawn from Theorem 2.2 (for instance by letting $p_{n}$ grows faster than $d_{n}$), or does it depend on the setting?  \n\n-**[Q3]**: In the examples discussed, the stochastic term $\\Sigma$ is only present around a fixed point of the drift part of the dynamics. Is there anything in the theorem preventing us from having $\\Sigma\\neq 0$ outside a fixed point?\n\n-**[Q4]**: What the authors mean by \"converging to a unstable fixed point\" (L268-L273, L316-L317)? I understand that the dynamics can get trapped in a neighbourhood of an unstable fixed points for some time if initialised close to it, but how can it converge to a unstable fixed point?  ", "limitations": "An inevitable limitation of the level of generality of this work is that it does not provide a constructive way to choose the summary statistics for a given problem of interest. This could be briefly commented in the text as well.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The investigation of one-pass SGD through the lenses of its summary statistics has a long history, which to my best knowledge dates back to the seminal work of Saad & Solla [41, 42], and was followed by an intense research activity since the mid-90s [BS, RB, CC]. These early works were inspired by ideas from statistical physics of disordered systems, where reducing the study of a high-dimensional random system to the study of low-dimensional equations for quantities that concentrate (known in that field as \"order parameters\") is quite natural. These early works have focused in the particular setting of two-layer neural networks (NNs) in a teacher-student setting (i.e. when the training data is generated by a two-layer NN itself) and in the critical scaling (adopting the author's terminology), but since have been enlarged to other data models, such as correlated Gaussians [20], Gaussian mixture classification [37], and to other scaling regimes [49]. I stress that because although [41] has been published in this very own venue (NeurIPS'95), closely related ideas have recently resurfaced in the mathematical machine learning literature, but missed this connection (e.g. [TV, 52]) - likely because they were independently rediscovered.\n\nIn this work, the authors take an important step forward by collecting and extending these specific settings into a general abstract framework for mapping the one-pass SGD dynamics of a probabilistic learning task into a SDE for the summary statistics. In particular, the characterisation of the stochastic correction in terms of a simple diffusion is an important contribution that goes beyond the literature above, and that allows to study how the dynamics escape from unstable fixed points.\n\nHowever, although most of the classic literature discussed above is cited in this manuscript, the connections are poorly acknowledged on the level of the discussion. For instance:\n\n- The trade-off between the \"population drift\" $f$ and the \"population corrector\" $g$ in the critical scaling already appeared in [41, 42] in the specific context of teacher-student two-layer networks. There, it was shown that the presence of $g$ gives rise to a fixed point that prevents the dynamics to reach perfect learning at linear time scales.\n\n- The cross-over between the \"critical\" and \"sub-critical\" regimes (when the correction $g=0$) with the relative scaling between the data dimension $d_{n}$, the number of parameters $p_{n}$ and the learning rate $\\delta_{n}$ was recently discussed in [49]. This should be reflected in the discussion in L63-L67 and L141-L144.\n\n- The SNR scaling for the existence transition of good classifiers for one-pass SGD on the Gaussian Mixture XOR classification task with two-layer NNs was discussed in [37].\n\nIn my opinion, the manuscript is also well-written and the conceptual thread is easy to follow. However, I believe the technical part in Section 2 and in particular Definition 2.1 is hard to parse for the general NeurIPS audience. It would be nice if the authors could add some intuition on what each of the conditions 1-3 mean, and maybe the simplest concrete example the author's can think of a localising and a non-localising sequence to help the reader develop an intuition.\n\nTo summarise, the strengths and weaknesses of the manuscript are:\n\n**Strengths**: The manuscript is well written and easy to follow. The theoretical contribution is a significant addition to the literature, setting a fairly general framework to understand the behaviour of one-pass SGD through tractable low-dimensional equations. The examples given are pertinent and useful to understand the theory.\n\n**Weaknesses**: The discussion is not well placed in the relevant literature. Some technical parts of the manuscript are hard to parse.\n\n\n**References** (numbered refs. are from the bibliography in the paper)\n\n[BS] M. Biehl and H. Schwarze, \"Learning by on-line gradient descent\", Journal of Physics A: Mathematical\nand General, vol. 28, no. 3, pp. 643–656, feb 1995.\n\n[RB] P. Riegler and M. Biehl. \"On-line backpropagation in two-layered neural networks\". Journal of\nPhysics A: Mathematical and General, 28(20), 1995.\n\n[CC] M. Copelli and N. Caticha, “On-line learning in the committee machine,” Journal of Physics A: Math-\nematical and General, vol. 28, no. 6, pp. 1615–1625, mar 1995.\n\n[TV] Y.S. Tan, R. Vershynin, \"Online Stochastic Gradient Descent with Arbitrary Initialization Solves Non-smooth, Non-convex Phase Retrieval\", arXiv: 1910.12837 [stat.ML]", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "gn5Ue-aCCKv", "reviewer": "Reviewer_ju2h", "summary": "This paper studies rigorous approximations of SGD dynamics. Assuming an SGD algorithm of the form\n$$ X_{n+1} = X_n - \\delta \\nabla L(X_n, Y_n), $$\nand a set of summary statistics $u: \\mathbb R^p \\to \\mathbb R^k$ for fixed $k$, the authors show that the dynamics of $u$ converge in probability\nto a stochastic differential equation of the form\n$$ du = (f(u) +  g(u))dt + \\Sigma(u)dW_t.$$\n\nThe authors also provide a phenomenology of the three terms in the RHS, wherein $f(u)$ corresponds to the gradient flow term, $g(u)$ is a deterministic correction term, and $\\Sigma$ encapsulates the variance of each SGD step. Depending on the chosen stepsize $\\delta$, those terms can be either influential or negligible in the associated SDE.\n\nThe remainder of the article is devoted to several examples. First, the authors study the use of SGD for matrix/tensor PCA:\n$$ Y = \\lambda x^{\\otimes k} + W, $$\nwhere $W$ is an i.i.d Gaussian $k$-tensor. With careful renormalizations of the dynamics, they map the SGD problem to an Ornstein-Uhlenbeck process, and show that the threshold for stability of this OU process matches the already known thresholds for hardness in terms of $\\lambda$.\n\nThe last two examples refer to Gaussian mixture classification; one is a simple binary task, while the other is the so-called XOR Gaussian mixture, which is not amenable to a simple linear separator. In both cases, the authors consider a two-layer neural network, and provide a full set of sufficient statistics to describe the dynamics. Those dynamics are characterized by a complex structure of fixed points (resp. 4 and 625 connected components), some of those being unstable; the authors then demonstrate how \"zooming in\" on the dynamics helps characterize the behaviour around those fixed points.\n\n", "questions": "- Can you provide additional explanations regarding the 3/32 probability of ending up at a stable fixed point in the XOR mixture ?\n- What would happen if, like in many real-world networks, the number of hidden units (and therefore, the number of sufficient statistics) starts to diverge ?", "limitations": "The limitations are adequately addressed.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This paper is the culmination of several previous works studying rigorous approximations of SGD dynamics (e.g. Tan and Vershynin, Ben Arous et al., Veiga et al. ...); it provides a very general result encompassing all previous works. The hypotheses are quite reasonable, and less restrictive than most usual restrictions (no uniform Lipschitz condition on $u$, only fourth moment conditions on $L$).\n\nThe provided examples illustrate quite well both the power and the limitations of this approach: on the one hand, the main theorem is applicable to a wide array of learning problems, and yields a very precise characterization as an SDE. On the other hand, those SDEs are highly non-trivial to study, even in the ballistic phase without any noise. Hence, there is still a lot of work to extract interesting insights from this theorem on a new learning task.\n\nThe main weakness of this paper is its terseness. The main theorem is introduced without giving much intuition about it, while the gist of the proof (disregarding the important aspects of martingale bounding) is based on a Taylor expansion of $u(X_{n+1})$ in powers of $\\delta$. This would give a much needed intuitive interpretation to the correction term and the diffusion matrix.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "c8Wb6XIMxsO", "reviewer": "Reviewer_uVoh", "summary": "This paper investigates the high-dimensional limit of SGD with constant step-size. The authors obtain an ODE/SDE type of limiting dynamics for the trajectory of a given finite set of summary statistics when the dimension of the data and the parameters go to infinity. The limiting dynamics is shown to be matching with the gradient flow for the population loss when the step-size is below the sub-critical scaling regime. When the step-size is in the critical scaling regime, it is shown that the limiting dynamics might possibly get a ballistic correction term and a diffusive term. In addition, the authors also develop a separate diffusive (SDE) limit for the rescaled summary statistics around a fixed-point of the unscaled ballistic dynamics in order to investigate the microscopic behavior of the dynamics around those fixed-points. The authors also demonstrate the applicability of their main theorem in several examples including tensor spike model, and classification with two-layer neural networks of binary and XOR-type Gaussian mixtures. ", "questions": "I would suggest the authors to include a notation section or a table in the main text or in the appendix. The introduction can be expanded more to motivate the subject better and can be divided into explicit subsections of prior works and contribution in order to clarifiy the limitations of the prior works and the advantages of your results better. It could also be helpful to mention the prior work and compare with your results in sections 3-5. Figures could be explained better either in the main text or in the captions. A summary illustration visualizing different limiting dynamics under different scaling regimes can included but the page limit might not allow that. \n\nI would also like to ask the following questions to the authors:\n1. What can you say about implicit bias/regularization of SGD in the critical scaling regime or in the microscopic regime with your main theorem? Can the population correction term be written as the gradient of another function which might imply that SGD in critical regime minimizes a new loss function with an additive bias term?\n2. What can your result tell about the over-parameterized regime where $p_n/d_n \\gg 1$?", "limitations": "Authors can discuss limitations of their results and future work in more detail. They mention the limiting dynamics of time-de4pendent rescaling as a future work in footnote 2.  ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The main result seems to apply to a general class of loss functions and data distributions compared to prior works.\n2. The theoretical claims are supported by experiments and explicit analysis of several benchmark models.\n\nWeaknesses:\n1. Presentation and the structure of the paper make it harder to follow.\n2. Contribution over the prior works is not very clear. \n\nOverall, I find this paper quite interesting and in good quality in terms of technical analysis. However, the material could be presented better for clarity. ", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept