Phase diagram of Stochastic Gradient Descent in high-dimensional two-layer neural networks
Authors: Rodrigo Veiga
Ludovic STEPHAN
Bruno Loureiro
Florent Krzakala
Lenka Zdeborova
Venue: NeurIPS 2022 Accept (2022)
Keywords: stochastic gradient descent
two-layer neural networks
overparametrization
gaussian inputs
statistical physics

Abstract: Despite the non-convex optimization landscape, over-parametrized shallow networks are able to achieve global convergence under gradient descent. The picture can be radically different for narrow networks, which tend to get stuck in badly-generalizing local minima. Here we investigate the cross-over between these two regimes in the high-dimensional setting, and in particular investigate the connection between the so-called mean-field/hydrodynamic regime and the seminal approach of Saad \& Solla. Focusing on the case of Gaussian data, we study the interplay between the learning rate, the time scale, and the number of hidden units in the high-dimensional dynamics of stochastic gradient descent (SGD). Our work builds on a deterministic description of SGD in high-dimensions from statistical physics, which we extend and for which we provide rigorous convergence rates.

Phase diagram of Stochastic Gradient Descent in high-dimensional two-layer neural networks
Rodrigo Veiga EPFL. Lausanne
Ludovic Stephan IdePHICS, EPFL, Lausanne
Bruno Loureiro IdePHICS, EPFL, Lausanne
IdePHICS,
IFUSP, USP, São Paulo
Florent Krzakala IdePHICS, EPFL, Lausanne
Lenka Zdeborová SPOC, EPFL, Lausanne
Abstract
Despite the non-convex optimization landscape, werrppaametrried shallow net- works able to achieve global convergence under gradient descent. The picture can be radically different for narrow networks, which tend to get stuck in badly- generalizing local minima. Here investigate the cross-over between these regimes in the high-dimensional setting, and in particular investigate the connection between the so-called mean- mean-field/hy field/hydrodynamii regime and the seminal approach of Saad & Solla. Focusing on the case of Gaussian data, we study the interplay between the learning rate, the time scale, and the number of hidden units the high-dimen nigh-dimensional dynamics of stochastic gradient descent (SGD). Our work builds on deterministic description of SGD in high-dimensions from statistical physics, which we extend and for which we provide rigorous convergence rates.
Descent-based algorithms such stochastic gradient descent (SGD) and its variants are the workhorse of modern machine learning. They are simple to implement, efficient to run and most importantly: they work well in practice. A detailed understanding of the performance of SGD is major topic in machine learning. Quite recently, significant progress was achieved the context of learning in shallow neural networks. In a series of works, it was shown that the optimisation of wide two-layer neural networks can be mapped to a convex problem the space of probability distributions over the weights [1-4]. This remarkable result implies global convergence of two-layer networks towards perfect learning provided that the number of hidden neurons large, the learning rate is sufficiently small and enough data is at disposition. This line of work is commonly referred to the mean-fiell or the hydrodynamic limit of neural networks. Mathematically. these works showed that one could describe the entire dynamics using partial differential equation (PDE) in d dimensions.
In a different, and older, line of work, one-pass SGD for two-layer neural networks with a finite number of hidden units, synthetic Gaussian input data and teacher-generated labels has been widely studied starting with the seminal work of [5]. These works consider the limit of high-dimensional data and show, in particular, that the stochastic process driven by gradient updates converge
set of p² deterministic ordinary differential equations (ODEs) as the input dimension d + oo and the learning rate is proportional to 1/d. The validity of these ODEs in this limit proven by [6]. However, the picture drawn from the analysis of these ODEs is slightly different from the mean mean-field/fe field/hydrodynamic picture: in this case SGD get stuck for long time in minima associated hidden units to the teacher hidden units, and even when it converges to specializing minima, it fails to perfectly learn (i.e. achieve zero population risk). In fact, in this analysis, the interplay between the limit of the learning rate going to zero and + oo appeared to be fundamental.
specialization of the
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

10"
Perfect learning
10"3
Bad learning
Platean
e6 1011
No ODEs
10"
00
10 19 10 10 10*
8
Figure 1: A summary of our contribution on SGD learning regimes for two-layer neural networks in the high-dimensional input layer limit d -> oc. Left: phase diagram according to Eqs. (1): Perfect learning region (K
0), Plateau line (x + 8 = 0), Bad learning region (-1/2 < K + 0) and No ODEs region (x 5 < 1/2). Right: A solution of the ODEs in all regions, with matching (2)). The time scaling is not uniform through the phase diagram: & = K ó on green and blue regimes and 0= 2(x+6)) on the orange region. The green curve decays power law to zero excess error
colors (using 0.301, p= 8, A: 4, Pre 6rev A = 10 (zz)
One should naturally wonder about the link between these two sets of works with, on the one hand a d-dimensional PDE (with large p), and on the other a p² -iimensiioaal ODE (with large d). this work we aim to build bridge between these two approaches for studying one-pass SGD.
Our starting point is the framework from [5], which build upon and expand to much broader range of choices of learning rate, time scales, and hidden layer width. This allows us to provide sharp characterisation of the performance of SGD for two-layer neural networks in high-dimensions. We show it depends on the precise way in which the limit is taken, and in particular on how the quantity of data, the hidden layer width, & the learning rate scales as d + OG. For different choices of scaling, we can observe scenarios such as perfect learning, imperfect learning with an unavoidable error, or even no learning all. As a consequence of our analysis, we provide phase diagram (see Fig.1) 1)' describing the possible scenarios arising in the high-dimensional setting. Our main contributions are as follow:
C1 We rigorously show that the dynamics of SGD can be captured by a set of deterministic ODEs, considerably extending the proof of [6] to accommodate for general time scalings defined an arbitrary learning rate, and a general range of hidden layer width. We provide much finer non- asymptotic guarantees which
crucial for our subsequent analysis.
C2 From the analysis of the ODEs, we derive a phase diagram of SGD for two-layer neural networks in the nigh-dimensional input layer limit - 0o. particular, scaling both the learning rate y and hidden layer width with the input dimension d as
x 5
p pxd"
identify different learning regimes which are summarized in Fig. 1(efft):
Perfect learning (green region, -õ): show that perfect learning (zero population risk) can be asymptotically achieved with n ~ + samples even for tasks with additive noise.
Plateau (blue line -δ): learning reaches plateau related to the noise strength. The point goes back to the classical work of [5].
Bad learning (orange region -1/2 < K < 0): here the noise dominates the learning process.
No ODEs (red region K + 8 1/2): the stochastic process associated to SGD is not guaranteed
to converge to a set deterministic ODEs. This region thus outside the scope of our analysis.
Relation to previous work Deterministic dynamical descriptions of one-pass stochastic gradient descent in high-dimensions have long tradition the statistical physics community, starting with
'Code: htgg:ttwb..omm//ddvviggg/a/i/iie

---

single- and two-layer neural networks with few hidden units [7-11]. The seminal work by [5] overcame previous limitations by constructing set of deterministic ODEs for two-layer networks with any finite number of hidden units, paving the way for a series of important contributions [12- 14, 6]. This line of work corresponds to the k. generalize this picture beyond fixed hidden layer size and learning rate.
case of Fig. (left). One of our goals is to
more recent line of work investigating the dynamics of SGD is the so-called mean-field limit [], 15, 2-4J, which connects the SGD dynamics of large-width two-layer neural networks to a diffusion equation in the hidden layer weight density. In particular, [15] provide non-asymptotic convergence bounds for sufficiently small learning rates, corresponding to the green region of Figure 1 (with p + oo). The mean-field approach computes the empirical distribution (in R²) of the hidden layer weights, while we focus on the macroscopic overlaps between the teacher and student weights.
1 Setting
Consider a supervised learning regression task. The data set is composed of n pairs (æ", yy )uE(x] E Rd+1 identically and independently sampled from P( y). The probability F(x) is assumed to be known and IP (z) is modelled by a two layer neural network called the teacher. Given a feature vector x" E R² the respective label y" E E is defined as the output of a network with k hidden units, fixed weights W E R and an activation function σ : R + R:
W x²
v" [a"", W*) + VA"" ; fa" W
Σ
Eo(A"") (2)
":: x" *=vvae E R as the r-th > 0 controls the strength of additive label noise: Ç" ~ IP((C) such that BG = 0 and ( 1.
with w. [W E R² as the r-th row of the matrix W and X w.
component of the teacher local field vector X E R The parameter A
Given a new sample for the respective label y. The error is quantified by a loss function L(y, f(x, O)), where + is an arbitrary set of parameters to be learned from data.
(x) outside the training data, the goal is to obtain an estimation f(x)
In this manuscript are interested in the problem of estimating W with another two-layer neural network with the same activation function, which we will refer to as the student. The student network has p hidden units and matrix of weights W E RPxd to be learned from the data. Given feature vector x ~ P((x) the student prediction for the respective label is given as
1 Σ p
w x
1
f(x. W)
(N))
p 3=1
where w, [] E Rdd is the j-th row of the matrix W and Xy a//a E IR is defined as the j-th component of the student local field vector X E R..
One-pass gradient descent- -
Typically, one minimizes the empirical risk over the full data set. Instead, learning with one-pass gradient descent minimizes directly the population risk:
R(W, W*) = Exy()())))) [c (f(x,, f"),f(z, f((ee))))]))))))))))))
(4)
Given single sample (x", y²) the weights are updated sequentially by the gradient descent rule:
w = wn - V L (9" [(x",("" w))
(5)
with v E [n]] and j E [p]. The parameter 7 > O is the learning rate. Despite being a simplification with respect to batch learning, one-pass gradient descent is an amenable surrogate for the theoretical analysis of non-convex optimization, since at each step the gradient is computed with a fresh data sample, which is equivalent performing SGD directly on the population risk. In particular. in this manuscript we assume realizability 2 k, and focus our analysis on the square loss L(y, 9) 3 (y- = leading to w w/ + 11
2 E" x" pvd
(6)
3

---

where
1 P P
E" _ o0;
Therefore, from the above expression we that to monitor the population risk along the learning dynamics it sufficient to track the joint distribution the local fields (2, X*). For Gaussian data P(x) (a 00 1), one replace the expectation E by Bx.A=~N(A.. oo.s) [+] and
fully describe the dynamics through the following sufficient statistics, known in the statistical physics literature as macroscopic variables:
Q" = Pa.y~P(a [2'2
W W
d
M Eu
W"W"
(8)
d
P E ((.a)
X
W W
with matrix elements, called order parameters in the statistical physics literature, denoted by ojt [Q",, mý. = jrr and Pre = [P],e
(r) The macroscopic state of the system the learning step is given by the overlup matrix Q2 E (R(++))x(ptt))
Q'
M"
n
(9)
M
P
and the population risk is completely determined by the macroscopic state:
(()) = EAA' EAA'~ BE-F(C) (
(10)
The training dynamics (5) defines a discrete-time stochastic process for the evolution of the overlap matrix { E (+++t))((++)) [nl]. with P fixed and Q" and M updated as:
2111011? ddd² E
u+1
O31
A + 14 X5) +
pd
learning
(11)
m
mîr
cV pd
learning
with v E [m]. [p]. E [k] and 8Y
= E" In what follows, we will make the concentration
assumption || = d: this will be justified in the proof of Theorem 2.1.
We emphasize in (11) the specific role played by each term in the right hand-side. The "learning" terms are the fundamental ones, that actually drive the learning of the teacher by the student. We show in Appendix C.3 that these "learning" terms are identical to those obtained in the gradient flow approximation of SGD, whose performance is the topic of many works [1-4]. Those are precisely the terms that draw the population risk towards zero. However, our setting there an additional variance term (so that this flow approximation is incomplete) that corresponds to the fluctuations of L(æ, W, W*) around its expected value R( R(W, W*). particular, this is where the effects of the noise ( can be felt. These terms were sometimes denoted as (T2) and (74) in [16]. We shall see that the additional "variance" term is the one responsible for the plateau in the critical (blue) region of Figure 1, while its contribution vanishes in the perfect learning (green) region.
Additionally, albeit our work particularizes Gaussian input data, we believe our conclusion, and the phase diagram discussed in Figure 1, to hold beyond this restricted case. Indeed, while the Gaussian assumption is crucial to reach a particular set of ODEs and their analytic expression, the approach can be applied to more complex data distribution, as long one can track the sufficient statistics required to have a closed set of equations. For instance, [17] obtained very similar equations for an arbitrary mixture of Gaussians that would obey the same scaling analysis as ours- while [18-20] proved that many complex distributions behave as Gaussians high-dimensional setting, including, e.g. realistic GAN-generate data. We thus expect our conclusions to be robust in this respect.

---

2 Main results
Although v seems to be the most natural time scaling in the high-dimensional limit d oc. if y and are allowed to vary with d the right-hand side (RHS) of Eqs. (11) can diverge and render the ODE approximation obsolete. Instead, for given time scaling öt, we can rewrite Eqs. (11)
μ+1
Lft
pd öt (E + 3 23) +
?
pa d 8 gr
öt
(12)
#+1 m
jr
mij
pd öt Ej X;
7
öt
In Theorem 2.1 we prove that as + oc, 52" converges to the solution of the ODE:
d (t) (S(())
(13)
dt
where gi (((++)x((++)) ((++ is the expected value the RHS of Eqs. (12), provided that this solution stays bounded. This enhances the result of [6] by providing convergence rates to the ODEs encompassing all scalings adopted hereafter:
Theorem 2.1 (Deterministic scaling limit of stochastic processes). Let E R be the continuous time horizon and öt (d) be a time scaling factor such that the following assumptions hold:
1. the time scaling öt satisfies for some constant c,
7 72 pd pdd
8t 2 C. max
(14)
2. the activation function is L-LLipchhiz
3. the function vi B (g(+t))((++b)
+ - ((+tt)x(p+))) is L'-Lipschitz.
Then, there exists a constant C > 0 (depending on c, L, L') such that for any S S [T/8t], the following inequality holds:
E
| Ω (v8t) S ee Cr log(p) vot.
(15)
Our proof is based on techniques introduced in [21] (namely, their Lemma 2) which studies a different problem with related proof techniques. The proof involves decomposing Ω 2²+1 as
n2+ S2" + (12" + (22 Ω2" 5tv((1")),
where the two first terms can be considered as a deterministic discrete process, and the last term is a martingale increment. The main challenge lies in showing that the martingale contribution stays bounded throughout the considered time period.
Although the method is similar to [6]. there are a number of differences between the two approaches. First, our proof fixes number of holes in [6]. in particular bounding 9", by sufficiently slowly diverging function of v. Additionally, the techniques used in this paper yield dependency in p that is nearly negligible, while the previous methods imply bounds that are much too coarse for our needs.
The function v can be computed explicitly for various choices of a, which allows check Assumption 3 directly. We provide in Appendix C the necessary computations for a (z) erf(æ:/ /2); those for the ReLU unit can be found in [22]. It can be checked that in the ReLU case, the function & is not Lipschitz around the matrices Ω satisfying
23
for some j * I. However, in every case we have weaker square-root-Lipschin property: there exists C' E R such that
(22) (12')) c ccllvv vs2i
for any 52, S2' Since the square root function is Lipschitz whenever the eigenvalues of Ω are bounded away from zero (see e.g. [23]), Assumption 3 is implied by the condition
22 2 elp-Ri
5

---

however, this assumption is much stronger, and becomes unrealistic in the specialization phase (as well as when p >> d).
Theorem 2.1 allows us to safely navigate through Fig. 1 by keeping track of convergence rates of the discrete process to a ODEs. The interplay between learning rate and hidden layer width defines the time scaling òt and the trade-off between the linear contribution on E and the quadratic one, playing central role on whether the network achieves perfect learning or not. Specifically, consider the following learning rate and hidden layer width scaling with d: 7 70/ dd P = pad where To E IR+ and Pn E N are constants. The exponent E R can be either greater or smaller than zero, while IR+, Replacing these scalings Eqs. (11), we find (where we have chosen To Po without loss of generality):
1+1
o3 d t1++ (E + d2+2(n+0) Ej
(16a)
learning
noise
m 2+1 m mir
+ E' 1.7
(16b)
learning
Since the distribution of the label noise P(() is such that [x))((y)s] 0, the linear contribution in E is noiseless in the high-dimensional limit d -> oo, and therefore we will refer to it as the learning term. The noise enters in the equations through the variance computed on the quadratic contribution E, Eg, which we will refer to the noise term; intuitively, it is a high-dimensional variance correction which hinders learning. In order to satisfy (14), we shall take
8t = max
(alticts' d(+2(x+6))
1
(17)
When K + 8 of O, this implies that either the learning term or the noise term scale like negative power of d, and is negligible with respect to the other term. is then easy to check that at a finite time horizon the resulting ODE behaves if the negligible term was not present. We refer to
Theorem B.1 the appendix for quantitative proof of this phenomenon. Let us now describe the different regimes depicted in Fig.
Blue line (plateau)- When and are scaled such that k = -6, Eqs. (16) converge to
dapt
1AN [8, À & + dia
EgsP(S) [
(18a)
dmjr = E.A.~ [E, A;
(18b)
dto
with 8lg 1/d. This regime is an extension of [5] for which x 5 0. The convergence rate observe for 8 is consistent
to the ODEs scales with 1/1 log(d), and the phenomenology with previous works studying the setting = = 0; namely the existence of an asymptotic plateau proportional to the noise level. For instance, the asymptotic population risk Rx is known to be proportional to > [6] when k = â = 0 and the dynamics is driven by rescaled version of Eqs. (18). Since the noise term does not vanish under this scaling, perfect learning zero population risk possible. There is always an asymptotic plateau related to the noise level A, and the learning rate
Green region (perfect learning)- If k > -8 we can define the time scaling 8t 1/d txd0 By Theorem 2.1, Eqs. (16) converge to the following deterministic set of ODEs:
dqpp
(AAA-A(A Ec-P(C) [E,&1]
PAA-~NAA [ A + E,A,+
(19a)
df
d=+6
dmjr (AAA [8
(19b)
dt,+6
at a rate proportional to d2 log(d), where we have highlighted that the noise term vanishes with (x++))) Hence, as long as x > -8 the noise does not play any role on the dynamics. This setting could be understood by taking effect learning rate Jeff d d==-6 on Roo & 7A, which

---

20
09
81
80
Am
AwB
107
Figure 2: Left: Population risk dynamics for == = 0 (Saad & Solla scaling) Po = 8, 4 Pre 6rs- Activation function: (x) = eef(f//v2 22). Data distribution: P(x)) Mo 1). Dots represent simulations (d 1000), while solid lines are obtained by integration of the ODEs given by Eqs. (18). Right: Simulations (d 1000, po = 8) for k = 0 comparing different choices of the exponent d. Noise level: 10-3, The final plateau proportional to the learning rate: Roo x 7A.
leads to zero population risk, i.e. perfect learning, in the high dimensional limit d + oc. We validate this claim by finite size analysis in the next section.
As discussed, the time scaling determines the number of data samples required to complete one learning step on the continuous scale. The bigger x + 6, the more attenuated the noise term, thus the closer perfect learning. The trade-off is that the bigger + 6, the larger the number of samples needed is, since n = rd + Given a realizable learning task, one would thus rather choose the parameters attain the perfect learning region, but being close possible to the plateau line for not increasing too much the needed number of samples. We remark that [15] provides an alternative deterministic approximation in this regime, with non-asymptotic bounds, whenever p 1; this is the so-called mean-field approximation, with known convergence guarantees [2].
Orange region (bad learning) We now step in the unusual situation where the learning rate grows faster with d than the hidden layer width: λ < -. In this case, by (17) the noise term dominates over the dynamics. Defining the time scaling 612(m45) = 1/d (x we have
dgg
(AAA E Bc [E, Et O (drtt),
(20a)
dt2(x48)
dmjr = O (d
(20b)
dt2(x+8)
According to Theorem 2.1 the convergence rate of Eqs. (16) to Eqs. (20) scales with (1(12++x+)) log( d). Therefore the existence of the noisy ODEs above circumscribed to the region 1/2 of the contributions of Eqs. (16) goes to zero, the worse is the convergence rate. In the present case, the more the learning term is attenuated, i.e. the more negative is K x Ó, the worse the dynamics is described by Eqs. (20). Although the weights are updated, the correlation between the teacher and the student weights parametrized by the overlap matrix M remains fixed on its initial value M which fixed point of the dynamics under this scaling. Unsurprisingly, this leads to poor generalization capacity.
0 and presents convergence trade-off absent in the other regimes: the faster
Red region (no ODEs) x 11/2, the stochastic process driven by the weight dynamics does not converge to deterministic ODEs under the assumptions of Theorem 2.1. We are then not able to state any claim about this regime.
Initialization and convergence There are two additional features worth commenting on the high- high-dimensional dynamies and its connection to the feaan-ielddydrodynammi approach, regarding initialization and the specialization transition. In the ODE approach we discuss here, we always observe a first plateau where the teacher-student overlaps are all the same. This means all the hidden layer neurons learned the same linear separator. At this point, the two-layer network is essentially linear. This called a unspecialized network [16, 5]. In fact, this a perfectly normal phenomenon, with few samples even the Bayes-optimal solution would be unspecialized [24]. Only by running

---

167
0% 1071
2m:30
10"1
1000
Figure 3: Networks parameters pa = 8, 4, Prs = 5rs- Activation function: o(x) = erf x / /v 2 Data distribution: P(æx) (x 1). Left: Population risk dynamics for K 0 and 5 1/2. Fixed noise A = 10-3 and varying d. Dots represent simulations. while the solid line is obtained by integration of the ODEs given by Eqs. (19). The data are compatible with the claim that
oo the curve converges to zero population risk. Right: Asymptotic population risk Re from simulations (dots) as function of d for different noise levels under the scaling 0 and 5 1/2. The fitted straight lines have slopes -0. -0.458. -0.494, -0 -0.497, for A = 10 9,103 7,100,, respectively.
81
1071
2-30
10'1
117
Figure 4: Networks parameters pa = 8, 4, Pre = -rs- Activation function: o(x) = erf( 2 /e///22 2 Data distribution: P(x))
0x 1). Left: Population risk dynamics for k 0 and 5 1/4. Fixed noise A = 10 1 and varying d. Dots represent simulations, while the solid line is obtained by integration of the ODEs given by Eqs. (19). The data are compatible with the claim that d the curve converges to zero population risk. Right: Asymptotic population risk Re from simulations (dots) function of d for different noise levels under the scaling 0 and 5 1/4. The fitted straight lines have slopes -0.201, -0.295, -0. -0.201, for A = 10 10-3 10- respectively.
the dynamics long enough the student hidden neurons start specialize, each of them learning different sub-function so that the two-layer network can learn the non-trivial teacher.
Let us make comments on this phenomenon: (i) while the "linear" learning in the unspecialized regime may remind the reader of the linear learning in the lazy regime [25, 26] of neural nets, the two phenomena are completely different. In lazy training, the learning is linear because weights change very little, so that the effective network is linear approximation of the initial one. Here, instead, the weights are changing considerably, but each hidden neuron learns essentially the same function. (ii) If the ODEs are initialized with weights uncorrelated with the teacher, then the unspecialized regime a fixed point of the ODEs: the student thus never specializes, any time. Strikingly, such condition arises as well in the analysis of mean-field equations (see e.g. Theorem 2 in [27] that discusses the need to have spread initial conditions with a non-zero overlap with the teacher) to guarantee global convergence. This raises the question about the precise dependence of the learning on the initialization condition in the high-dimensional regime, where a random start gets vanishing (1/vd) overlap. This challenging problem that only recently has been studied (though in a simpler setting) in [28-30] who showed it yields an additional log(d) time-dependence. Generalizing these results for high-dimensional two-layer nets an open question which we leave for future work.
Discussion, special cases, and simulations
To illustrate the phase diagram of Fig. I, present several special cases for which we can perform simulations or numerically solve the set of ODEs. Henceforth, we take oxx) = crf(x// v2).

---

for which the expectations of the ODEs and of the population risk, Eq. (10), can be calculated analytically [5]. The explicit expressions are presented in Appendix C. Teacher weights are such that Pro = ō.s- The initial student weights are chosen such that the dimension d can be varied without changing the initial conditions Q". M P and consequently the initial population risk Ro. A detailed discussion can be found in Appendix D.
Saad & Solla scaling K = 0- We start by recalling the well-known setting characterized by the point K = 8 = 0. The convergence of the stochastic process for fixed learning rate and hidden layer width to Eqs. (18) was first obtained heuristically by [5]. In Fig. 2 (left) we recall this classical result by plotting the population risk dynamics for different noise levels. Dots represent simulations, while solid lines are obtained by integration of the ODEs, Eqs. (18).
Learning is characterized by two phases after the initial decay. The first is the unspecialized plateau where all the teacher-studes overlaps are approximately the same: mjr R m. Waiting long enough, the dynamics reaches the specialization phase, where the student neurons start to specialize, i.e., their overlaps with one of the teacher neurons increase and consequently the population risk decreases. This specialization is discussed extensively [5]. A = O, the population risk goes asymptotically to zero. Instead, if A + O, the specialization phase presents second plateau related to the noise A.
The asymptotic population risk Reo related to the 2nd plateau is proportional to > [6] in the high-dimensional limit d + oo with p finite. As mentioned in the previous section, the expectation over & & in Eq. (18a) prevents one from obtaining zero population risk for noisy teacher.
Perfect learning for x 0 - In this section we study the line k= 0 with d 0 of Fig. 1, for which Eqs. (19) with x = 0 hold. We show that perfect learning can be asymptotically achieved in the realizable setting for any finite hidden layer width p Po. Keeping 5 and A fixed, we have done simulations increasing the input layer dimension d. In Fig. 3 (left) we set 8 = 1/2, A = 10 and vary the input layer dimension. The bigger d is, the closer we are the ODE-derived noiseless result.
Gathering the asymptotic population risk from simulations for varying d and A we perform a finite- size analysis to study the dependence of Reo with d. This shows that the noise term goes to zero under this setting. In Fig. 3 (right) we plot Rx versus d from simulations (dots) for different noise levels. We fit lines under the log-log scale showing that Rx d as expected. Fig. 4 draws the same conclusion for d = 1/4.
As already stated, the interplay between the exponents directly affects the time scale. We end this subsection by graphically illustrating this fact through simulations. Setting the noise to A 10-3 we compare the cases 8 = 0, 1/4, 3/8, 1/2 in Fig. 2 (right). All simulations are rendered on the scale 8Lg= = 1/d to illustrate the trade-off between asymptotic performance and training time.
Bad learning for K. 0- We now quickly discuss the uncommon case of growing with d within the orange region. In Fig. 5 (left) we compare simulations varying d with the solution of the ODEs given by Eqs. (20). Both lead to poor results compared the green & blue regions. Moreover, this regime present strong finite-size effects, making it harder to observe the asymptotic ODEs at small sizes. However, the trend increases very clear from the simulations. As discussed in Section 2, the more the learning term is attenuated on the ODEs, the worse they describe the dynamics.
Large hidden layer: K > 0 Finishing our voyage through Fig. with examples, we briefly discuss the case where both input and hidden layer widths are large. Although Theorem 2.1 provides non-asymptotic guarantees for K > O, the number of coupled ODEs grows quadratically with p. making the task of solving them rather challenging. Thus, we present simulations that illustrate the regions of Fig. 1. Fixing d = 100 we show in Fig. (right) learning curves for different values of k and 8. The colors are chosen to match their respective regions in the phase diagram.
Due to the relatively small sizes used in Fig. (right), the green dots seem decrease towards perfect learning, even when 5 < 0. provided that k is large enough, as is predicted by the phase diagram in Fig. 1. Moreover, since d is not large enough, when the parameters are within the orange region the finite-size effects actually dominates, similarly to Fig. 5. The learning contribution still plays role and the asymptotic population risk is similar to the case == 0. Within the red region, which out of scope of our theory, the simulation gets stuck plateau with larger population risk.

---

10"1 adc
30
A
444000
38
10
Figure 5: Left: Population risk dynamics for y and d 3/8. Fixed noise A 10-3 and varying d on simulations. Dots represent simulations, while the solid line is obtained by integration of the ODEs given by Eqs. (20). Right: Simulations across different regions Fig. 1. Networks parameters d 100, p d', (x) = erf( /z/ 2). Data distribution: P(x)) = (2 0, 1). Time scaling: = 5 for green and blue and 1 2(x + 8) for orange. The colors match Fig. 1.
d k = 4, Prs 87s- Noise: A 10-3. Activation function:
4 Conclusion
Building up on classical statistical physics approaches and extending them to a broad range of learning rate, time scales, and hidden layer width, rendered sharp characterisation of the performance SGD for two-layer neural networks in high-dimensions. Our phase diagram describes the possible learning scenarios, characterizing learning regimes which had not been addressed by previous classical works using ODEs. Crucially, our key conclusions do not rely on an explicit solution, our theory allows the characterization of the learning dynamics without solving the system of ODEs. The introduction of scaling factors is non trivial and has deep implications. Our generalized description enlightens the trade-off between learning and hidden layer width, which has also been crucial in the mean-field theories.
Acknowledgements
We thank Gérard Ben Arous, Lenaic Chizat, Maria Refinetti and Sebastian Goldt for discussions. We acknowledge funding from the ERC under the European Union's 's Horizon 2020 Research and Innovation Program Grant Agreement 714608- SMiL.e, and by the Swiss National Science Foundation grant SNFS OperaGOST, 200021 1_1329000 RV was partially financed by the Coordenação de Aperfeigoamento de Pessoal de Nível Superior Brasil (CAPES)- Finance Code 001. RV grateful to EPFL and IdePHICS lab for their generous hospitality during the realization of this project.
10

---

References
[1] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):
2018. ISSN 0027-8424. doi:10. 1ett100.173/pmaa 1806579115.
[2] Lénaïc Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Pro- cessing Systems, volume 31. Curran Associates, Inc., 2018
[3] Grant Rotskoff and Eric Vanden- Vanden-Eijnden Eijnden. Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, (5(9):1899-1935 2022.
[4] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications, 130(3):820-1152, 2020. ISSN 0304-4149, doi:httss://oii.org 1e..1111116jj.. 101.1100.......000 2019 .00.003..
[5] David Saad and Sara A. Solla. On-line learning in soft committee machines. Phys. Rev. E, 52: 4225 4243, Oct 1995. 1oi: 103/Phhyeeee...522255....
[6] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. H. Wallach, H. Larochelle, A. Beygelzimer, d' d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[7] W Kinzel and P Ruján. Improving a network generalization ability by selecting examples. Europhysics Letters (EPL), 13(5):473-477, nov 1990. doi:10. 100.02229990992555/53355550
(8] O Kinouchi and N Caticha. Optimal generalization in perceptions. Journal of Physics A: Mathematical and General, 25(23):6243-6255) dec 1992. 00000000000000555477223333020000
[9] M Copelli and N Caticha. On-line learning in the committee machine, Journal of Physics A: Mathematical and General, 1615)161,, mar 1995. 1100000888/0305447722666666666
[10] M Bichl and H Schwarze. Learning by on-line gradient descent. Journal of Physics A: Mathematical and General, 28(3):643-656, feb 1995. doi: 01f::000088/03555447722833333333
[11] P Riegler and M Biehl. On-line backpropagation in two-layered neural networks. Journal of Physics A: Mathematical and General, 28020)) oct 1995. 10::1001088/083
4470/28/20/002.
[12] David Saad and Sara Solla. Dynamics of on-line gradient descent learning for multilayer neural networks. In D. Touretzky, M. C. Mozer, and M. Hasselmo, editors, Advances in Neural Information Processing Systems, volume 8. MIT Press, 1996.
[13] Renato Vicente. Osame Kinouchi, and Nestor Caticha. Statistical mechanics of online learning of drifting concepts: A variational approach. Machine learning, 32(2):179-201, 1998. ISSN 08885-6155. doi: 10. 1023/A:100742873199
[14] David Saad, editor. On-Line Learning in Neural Networks. Publications of the New- ton Institute. Cambridge University Press, Cambridge, 1999. ISBN 9780521652636.
10.1100..11/BBB 1001000///978011159999
[15] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-fiee bounds and kernel limit. 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 2388-2464. PMLR, 25-28 Jun 2019.
[16) David Saad and Sara A. Solla. Exact solution for on-line learning in multilayer neural networks. Phys. Rev. Lett., 74:4337-4340, May 1995. doi: 10. 2001:00110//PyysReeee7144..3333

---

[17] Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborova. Classifying high- dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9936-8947, PMLR, 18-24 Jul 2021.
[18] Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Mare Mézard, and Lenka Zdeborová. The gaussian equivalence of generative models for learning with two-layer neural networks. In Proceedings of Machine Learning Research, volume 145, pages 1-46. 2nd Annual Conference on Mathematical and Scientific Machine Learning, 2021.
[19] Hong Hu and Yue M. Lu. Universality laws for high-dimensional learning with random features. IEEE Transactions on Information Theory, pages 1-1, 2022. 109.111///TT222222.33.3 126777
[20] Andrea Montanari and Basil N. Saeed. Universality of empirical risk minimization. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Reseurch, pages 4310-4312. PMLR, 02-05 Jul 2022.
[21] Chuang Wang, Yonina C. Eldar, and Yue M. Lu. Subspace estimation from incomplete observa- tions: A high- nigh-dimensional analysis. IEEE Journal of Selected Topics in Signal Processing, 12 (6):1240-1252, 2018. doi: 10w.11/0///STS.20188887..0
[22] Yuki Yoshida, Ryo Karakida, Masato Okada, and Shun-ichi Amari Statistical Mechani- cal Analysis of Online Learning with Weight Normalization in Single Layer Perceptron. Journal of the Physical Society of Japan, 86(4):044002, April 2017. ISSN 0031-9015. doi: (0..5566/JPS1.86.4444
[23] P. Del Moral and A. Niclas. A taylor expansion of the square root matrix function. Jour- nal of Mathematical Analysis and Applications, 465(1):259 266, 2018. ISSN 0022-247X.
doi:httss:/doi.com org/10 1...01160..jmaa..e 200...008.00...........100 20...05.005.
[24] Benjamin Aubin, Antoine Maillard, Jean Barbier, Florent Krzakala, Nicolas Macris, and Lenka Zdeborová. The committee machine: computational to statistical gaps in learning a two-layers neural network. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):124023, dec 2019.
[25] Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
[26] Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. In H. Wallach, H. Larochelle, A. Beygelzimer, F. AAlcé--Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates. Inc., 2019.
[27] Francis Bach and Lénaic Chizat. Gradient descent on infinitely wide neural networks: Global convergence and generalization. arXiv preprint arXiv:21 10.08084, 2021.
[28) Yan Shuo Tan and Roman Vershynin. Phase retrieval via randomized kaczmarz: theoretical guarantees. Information and Inference: A Journal of the IMA, 8(1):97-233, 2019.
[29] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high- nigh-dimensional inference. Journal of Machine Learning Research, 22(106):1-51, 2021.
[30] Gérard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Algorithmic thresholds for tensor PCA. The Annals of Probability, 48(4):2052 2087, 2020. doi: 101100.214/99--0P1415
[31] Chuang Wang, Hong Hu, and Yue Lu. A solvable high-dimensional model of gan. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Bue, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
12


Reviews:
[{"review_id": "p4h6mIrHdtI", "reviewer": "Reviewer_dAar", "summary": "This paper considers the behavior of online SGD on a two-layer neural network in high-dimension, for different scaling of the number of neurons and step sizes. The authors consider a student-teacher scenario and gaussian data. In this setting, the test error is fully characterized by the covariance matrix of the student and teacher weights. When the time discretization is small enough, and for finite (continuous) time horizon, they show that the SGD dynamics concentrates on a deterministic ODE (with non-asymptotic bounds). The number of neurons and the step size scale polynomially with the dimension, and the authors identify four different regimes, depending on the ratio of the step size over the number of samples: perfect learning (error at the end vanishes when $d \\to \\infty$), plateau (proportional to the noise level), bad learning (dominated by noise) and no ODE. Plenty of numerical simulations are provided to illustrate their theoretical results.\n", "questions": "I list below a few comments, questions and concerns.\n\nIn the introduction, the link between Saas & Solla ODE and the mean-field PDE is mentioned. However, this is not discussed further in the main text. It would be interesting to define the mean-field dynamics in details: in the mean-field regime, the number of neurons p does not need to grow with the dimension d (as written line 146), but only need to be large enough (for a fixed time horizon). A propagation-of-error analysis show that $\\delta << p d  / (d + \\log (p))$ is enough for SGD to concentrate on the population gradient flow (perfect learning). This deserves further discussion. More generally, the readers might be more familiar with the NTK scaling or abc-parametrization (Greg Yang 2021) and it would be helpful to compare the scaling in this work to those lines of work.\n\nIn the evolution equation of the overlap matrix (between line 112 and 113), there is a multiplicative factor $\\| x_\\nu \\|_2^2 /d$ that is missing in the noise part. While it is true that this quantity concentrates on $1$, this should be analyzed (please propagate in the proofs if need be, I don’t expect it to be difficult or change anything to the story). If I understand correctly, the decomposition in learning and noise parts is related (but not equal) to a bias variance decomposition on the label noise (here the two parts have different scalings, hence this decomposition). However, when label noise is 0, the noise term has a fixed point for perfect learning. I am unclear about its interpretation for 0 noise. Maybe some comments about this would be helpful.\n\nMore generally, what is the connection between the ODE and the different regimes of SGD?  For classifying the different regimes, only the scaling of the learning and noise parts are compared. What does the ODE perspective provide? (It is difficult to analyze and simulate) Can it be shown that the risk converges to a value proportional to the noise level for parameters $\\kappa \\neq 0$ ? \n\nContrary to the discussion below Theorem 2.1, I do not think Assumption 3 is mild: for example, it cannot be true if $p > d$.  \n\nEven though this is not the focus of the paper, is there any sense when global convergence occurs in finite time horizon? Should the random start have non-vanishing gradient in the direction of the teacher?\n\nLine 206: *leads", "limitations": "Yes.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Understanding the behavior of SGD in neural networks is central to a theory of deep learning. However only a few results exist on the subject. While this paper avoid the issue of global convergence of the ODE, it provides a simple characterization of different regimes based on the label noise and the scaling of the width and step size. I think this is an interesting work, which the NeurIPS community will enjoy, and I recommend to accept the paper. ", "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": "do50bkUWZym", "reviewer": "Reviewer_StPv", "summary": "The authors rigorously analyze the deterministic dynamical descriptions of one-pass SGD of a two-layer neural network with the interplay between the learning rate $\\gamma\\propto d^{-\\delta}$, the time scale, and the number of hidden units $p\\propto d^\\kappa$ on a $d$-dimensional Gaussian data as $d\\rightarrow \\infty$.", "questions": "- L17: Why is it a convex problem?\n- What does it exactly mean by \"the one-pass SGD\"?\n- L202: What does it mean by \"a convergence trade-off\"?\n", "limitations": "Please provide further limitations when introducing the specific setting considered in the paper.", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strength\n- very rigorous approach\n- much more generalized settings compared to previous works (e.g., Saad & Solla 1995).\n\n### Weaknesses\n- The only weakness I found is the setting the paper considered. \nI'm not sure how the analysis on (i) the specific setting in the paper can help understand the dynamical properties of (ii) a common neural network training.\nIt would be better for a broad range of community if the authors could provide further implications between the two settings:\n- (i) convex, two-layer, global convergence, with teacher-generated labels, limit of high-dimensional data $d\\rightarrow \\infty$, one-pass, Gaussian data\n- (ii) non-convex, multi-layer, no global convergence, without teacher, finite high-dimensional data $d<\\infty$, (not one-pass?), general non-Gaussian data\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "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": "b73C_qkJa3S", "reviewer": "Reviewer_E938", "summary": "This paper studies the behavior of two-layer neural networks trained by online SGD in the high-dimension limit. In a teacher-student setting with Gaussian training data, an ODE system is shown to be the high-dimension limit of the dynamics of an \"overlap matrix\" that fully governs the behavior of the population risk. With the limit ODE system, different cases regarding the scaling of learning rate and network width with respect to the dimension is studied, and a phase diagram is plotted with the two scalings. The phase diagram consists of four regions---perfect learning, plateau, bad learning, and no ODE. The first three regions are determined by the relative scale of the noise and the learning terms in the dynamics. Although no explicit characterization is made for the dynamics or the limit of the loss, numerical experiments show that the three regions in the phase diagram indeed lead to different behaviors. ", "questions": "1. In the introduction the authors mentioned to connect the finite dimensional and infinite width case (the mean-field regime) with the infinite dimensional and finite width case. Can the authors make more discussion on the implication of theoretical results in the paper on the mean-field regime? It seems all the analysis in the paper treat the infinite-dimensional case, while in the mean-field analysis the dimension is fixed to be finite. \n\n2. The neural network models studied in the paper do not have output parameters after the activation function, which is quit different from a usual two-layer neural networks. Does the theory apply to two-layer neural networks with parameters in both layers?", "limitations": "The authors have honestly discussed the limitations of the work. ", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality: The ideas of studying high-dimensional limit of neural network training and using \"overlap matrix\" to characterize the behavior come from previous works. However, analysis in previous works is significantly improved to give the results in the current paper. Studying the impact of the two scaling factors and plotting the phase diagram are also original to this work. These results can deepen our understanding to the training of two-layer neural networks. \n\nQuality and Clarity: The paper is very well written. All the mathematics and discussions are clear and easy to read.\n\nSignificance: The results are theoretically significant because they captures the behavior of the training of neural networks in an extreme case---the high dimensional case. The training process of neural networks is very hard to study, due to its nonlinear nature. Therefore, understanding an extreme case that might not be realistic is still helpful and important. ", "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": "PNHc5fiz3w", "reviewer": "Reviewer_Q7z1", "summary": "The manuscript proves an ODE-limit (i.e. fluid limit) for the training trajectories of Soft Committee machines in the Teacher/Student setting and when the dimension of the inputs $d \\to \\infty$. Interestingly, the paper considers the regime when the number of neurons $p \\geq 1$ in the hidden layer and the learning rate $\\gamma>0$ can scale quite freely with the dimension $d$ of the inputs.", "questions": "1. The function $\\psi$ that is defined as the expected value of the RHS of Eqs. (12) is, I believe, depending on $(p,d,\\delta t)$, although I think I understand that there is a well-defined limit under appropriate scaling of $(p,d,\\delta t)$. If that is the case, at the cost of slightly more cumbersome notations, it would be more appropriate to reflect this dependency.\n\n2. There is quite a large literature on fluid-limit / ODE limits for analyzing randomized algorithms / Markov chain. Standard references are the book by Ethier & Kurtz, as well as the article \"Differential equation approximations for Markov chains\" by Darling & Norris, and many other places. I think that it would be worth to slightly expand the literature review on these fluid limits.\n\n3. Once Equation (12) is obtained, at an heuristic level at least, it is very clear what the scaling of $(p,d,\\delta t)$ needs to be in order to observe an ODE limit. Indeed, as soon as  $\\delta t \\to 0$ with $E( \\text{(markov chain increments)} ) \\asymp \\delta t$ and $E( \\text{(markov chain increments)}^2 ) \\ll \\delta t$ one expects a ODE limit when trajectories are appropriately time-rescaled: this already heuristically gives more or less all the results discussed after the statement of the main theorem. Indeed, the proven theorem is very interesting and *much* stronger, but I feel it would be worth mentioning that the conclusions are, in some sense, not entirely surprising once we have Eq. (12).\n\n4. The important discussion of the \"learning\" and \"noise\" terms in Eq. (16) [i.e. very bottom of page 5] could be expanded further since it is quite crucial to the existence of plateaus / etc... For example, why is this the case in general that (1) if the noise term vanishes  perfect learning is achievable (2) if the noise term remains as in the $\\kappa+\\delta = 0$ then perfect learning is not.\n\n5. All the weights of the \"second layer\" are fixed (and equal): it was not clear to me if it was a technical detail to improve readability (which is entirely fine!), and if the proof/approach could be relatively easily extended to the more general case [eg. as in article (6)]?\n\n**[responses by the authors acknowledged (4/8/2022)]**\nThank you for these clarifications.\n\n", "limitations": "NA.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is well written, with a reasonable literature review and a good balance between rigor and intuition. Furthermore, to the best of my knowledge (although I am *not* an expert in this area), the proven results are new. These results offer a comprehensive understanding of the generalization properties of these soft committee machines in the online regime (i.e. infinite training data) and as $d \\to \\infty$.\n\nAs weaknesses, I have only a few relatively minor remarks that are in next section.", "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": ["I don’t know"]}, "extra_text": {}}, {"review_id": "D8IWFnwekc", "reviewer": "Reviewer_8rq6", "summary": "This paper present an extension of results obtained in the mid 90's by Saad and Sola, who showed that stochastic gradient dynamic in the context of student/teacher learning of a one hidden layer network for gaussian data can be replaced by ordinary differential equations on the weights *when the dimension of the input is very large*.\nThey extend this work to number of hidden units that scales with the input dimension and to learning rate that can also scale with the dimension.\nThe authors aim at bridging the Saad and Sola approach with the mean-field limit of a large number of hidden units which is known to overcome the non-convexity of narrow networks for fixed dimension.\n\n", "questions": "- The authors use a scaling of the learning rate $\\gamma$ like $\\gamma_0/d^\\delta$ with a constant $\\gamma_0$ and a parameter $\\delta$ that can be either positive or negative. A negative parameter $\\delta$ seems a bit unusual. It would be interesting to comment on the practical relevance of this case. I understand that it highlights the fact that one can choose a negative $\\delta$ when the scaling of hidden units is polynomial in the dimension. I would like the authors to give more insight on the practical relevance of their result in this situation.\n\n- How could the authors design an experiment to assess that their results is not limited to the gaussian setting? This would give more credit to their work.", "limitations": "Not applicable.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- Scaling of learning rate and number of hidden units in terms of the dimension is a real improvement over the Saad and Sola setting.\n- the paper is quite well written.\n- So far I checked, the proofs are ok.\n\nWeaknesses:\n\n- The proposed study is limited to the realm of gaussian data and the authors believe the result is applicable more generally but no experiments is made to assess these beliefs. \n- It is not completely clear what are the consequences of these results on practical instances of the problem. Since the scalings are formulated with respect to the dimension, it would be interesting for the reader to make more concrete the practical consequences.", "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": "7EgpxQWmn9", "reviewer": "Reviewer_JBZo", "summary": "This work draws a phase diagram of SGD for two-layer neural networks in high-dimensional settings by connecting the mean-field regime with an approach by [Saad & Solla]. Specifically, this diagram categorizes learning dynamics depending on the order of network width and learning rate into four types: perfect learning, plateau, bad learning, and no ODEs. ", "questions": "My major concern is the third condition in Theorem 2.1. To rigorously justify the claim of the paper, this assumption should hold with a long time horizon, but it seems difficult. See the comment in [Strengths and weaknesses] for detail. \nIf my understanding is wrong, I would appreciate it if you could correct me.\n", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths** \n\nIt is well known that the choice of hyperparameters (learning rate, network width, initialization scale, etc.) significantly affects the generalization performance of neural networks trained by SGD. However, theoretical understanding of this phenomenon has been limited (e.g., theory of mean-field regime and neural tangent kernel.) In this sense, this work is interesting because the authors aim to characterize learning behavior more precisely and derive a phase diagram that matches an empirical observation. \n\nThe proof is a non-trivial extension of the previous work [Goldt et al. (2019)]. In particular, this work proved that a trajectory of population risk achieved by SGD converges to a deterministic ODE and investigated an obtained error by relating a trajectory with the limit of ODE. This approach is potentially important to perfectly understand the learning behavior of two-layer neural networks. \n\n**Weaknesses**\n\nHowever, I have a concern about Assumption 3 made in Theorem 2.1 which assumes $\\overline{\\Omega}(t)$ is strictly positive until a time horizon $\\tau$. And it is shown that time $\\tau$ also terminates the approximation guarantee Eq. (15). Based on these, we can see that unless Assumption 3 holds with $\\tau=\\infty$, we cannot relate the learning dynamics of SGD with the limit of ODE of $\\overline{\\Omega}$. Moreover, I think it is difficult to Assumption 3 holds with a long time horizon because this condition implicitly requires linear independence of features in student and teacher models but basically, $W^\\nu$ will degenerate by the convergence to the teacher model.\n\n\n**Other Comments**\nThe following paper is closely related to this work. [Yang and Hu (2021)] also categorizes hyperparameter settings of the learning rate, network width, and initialization scale into a learning regime that exhibits feature learning property, and derives preferable settings. \n\n[Yang and Hu (2021)] Greg Yang and Edward J. Hu. Feature Learning in Infinite-Width Neural Networks. ICML, 2021.", "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