Phase diagram of Stochastic Gradient Descent in high-dimensional two-layer neural networks

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.

Reviews and author responses:

 We thank the reviewer for the comments and for supporting the acceptance of our work.

 We thank again the reviewer for the constructive comments. Below we address the remarks.

>*The propagation of error analysis was done in "Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit", Mei et al. (2019) (note that one need to be careful when comparing the two papers, and need to use the following correspondence $\theta_j = w_j/ \sqrt{d}$ between the two papers...). Maybe it would be good to comment on this, as this implies that for any $\delta,\kappa$ scaling that makes this non asymptotic bounds vanish, the dynamics is well approximated by the population gradient flow with infinite width $p = \infty$ for number of iterations $pd/\gamma = d^{1 + \kappa + \delta}$ (it needs to be double checked, but it seems that the entire perfect learning region correspond to dynamics that are well approximated by the distribution dynamics PDE/mean-field limit).*

We thank the reviewer very much for these comments. We will cross-check in detail the scalings with the non-asymptotic bounds in Mei et al 2019 and discuss the relation in the revised manuscript.

>*About the assumption on $\Omega_t$, I was making an even simpler observation: $\Omega_t$ is the covariance of $p+k$ vectors in $d$ dimension, so $\Omega_t$ cannot be full rank whenever $p+k > d$ (it is at most rank $d$, whatever noise is injected). In particular, if $\kappa >1$, this assumption cannot be verified. I can believe that $\Omega_t$ remains full rank during the dynamics when $\kappa<1$, but the paper also considers $\kappa >1$.*

This is indeed a very good point, we thank the reviewer for their clarification of the question. The assumption should indeed be stated more carefully: for $\kappa>1$ the rank needs to be d. Assumption $3$ is only needed to ensure that the square root function is Lipschitz along the trajectory: When $p+k>d$, however, the dynamics is confined to a $d$-dimensional manifold (the set of matrices of the form ${\bf X}^\top {\bf X} $ for ${\bf X}$ of size $(p+k) \times d$). In this case the square root function should be Lipschitz on *this manifold*. Our proof goes trough also for $\kappa >1$ but we need (and will of course) to restate the assumption more carefully in the revised manuscript.

 Thank you for the additional question

A distinctive property of the NTK regime is that the learning rate is so small that the weights of the first layer move only infinitesimally away from their initial value. In the scaling/regime we consider the weights move by ${\cal O}(1)$ amount. In our scaling the NTK would correspond to an extremely small learning rate and a very short time. 

Therefore, our results are not contradicting the ones cited by the reviewer. The learning rate needs to be in an intermediate regime large enough so that features are learned yet not too large so that GD behaves reasonably.  When the learning rate is really large, learning always fails. In fact, this is well seen in practice, where it is common to try choosing a large learning rate (for speed) but not too large - otherwise, GD might take such a big step that the training loss does not even decrease. 

We will comment on the NTK regime in more detail in the discussion of related work to clarify this point.

 I am satisfied with it.

 Thank you for the answers.
> We will add a paragraph to the conclusions clarifying the general positioning of our work. Our setting is a shallow neural network in **the feature learning regime** - the same regime considered in the works on **the mean-field limit**. As argued in a number of works in this literature, we believe most of the state-of-the-art neural nets operate in such **a feature rich regime**. This is in stark contrast to **the NTK regime**, which ultimately reduces to **a convex kernel method**. Hence, NTK is considerably simpler from both the optimization and analysis point of view, and **fails to capture feature learning**, which is one of the key aspects of state-of-the-art neural networks.

I have a question (you may not answer the question as we are almost at the end of the discussion period): usual neural networks generalize better when trained with a large learning rate $\gamma$.
For example, Lewkowycz et al. [*] said that (in abstract)
> "In the small learning rate phase, training can be understood using the existing theory of infinitely wide neural networks. At large learning rates the model captures qualitatively distinct phenomena, including the convergence of gradient descent dynamics to flatter minima."

I think the two phase corresponds to **kernel (NTK) regime** (small $\gamma$) and **feature rich regime** (large $\gamma$), but this paper suggests somewhat opposite - **perfect learning** (small $\gamma$) and **bad learning** (large $\gamma$). Why do you think a smaller $\gamma$ with a larger $\delta$ leads to "perfect learning" in your settings?


[*] Lewkowycz et al., The large learning rate phase of deep learning: the catapult mechanism (https://arxiv.org/abs/2003.02218)






 Thank you for the answer.

The propagation of error analysis was done in "Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit", Mei et al. (2019) (note that one need to be careful when comparing the two papers, and need to use the following correspondence $\theta_j = w_j/ \sqrt{d}$ between the two papers...). Maybe it would be good to comment on this, as this implies that for any $\delta,\kappa$ scaling that makes this non asymptotic bounds vanish, the dynamics is well approximated by the population gradient flow with infinite width $p = \infty$ for number of iterations $pd/\gamma = d^{1 + \kappa + \delta}$ (it needs to be double checked, but it seems that the entire perfect learning region correspond to dynamics that are well approximated by the distribution dynamics PDE/mean-field limit)

About the assumption on $\Omega_t$, I was making an even simpler observation: $\Omega_t$ is the covariance of $p+k$ vectors in $d$ dimension, so $\Omega_t$ cannot be full rank whenever $p+k > d$ (it is at most rank $d$, whatever noise is injected). In particular, if $\kappa >1$, this assumption cannot be satisfied. I can believe that $\Omega_t$ remains full rank during the dynamics when $\kappa<1$, but the paper also considers $\kappa >1$.

 I will keep my score.


 The authors have addressed the concern well. 
I was concerned about the degeneration of $\overline{\Omega}$, but I understood that a (tiny) mismatch between parameters made by small $\Delta>0$ will prevent the degeneration. Thanks for the clarification. I would like to raise the score.

 We thank the reviewer for the follow-up questions and for supporting the acceptance of our work.

> *Is there anything in common between properties of neural networks and the observations in the paper? There are many simple models, but why do you consider the settings in the paper. Is it just for analytic investigation? Does it have similar properties with neural networks? Is there anything similar to the convergence trade-off observed in neural networks? For example, the neural tangent kernel (NTK) has similar optimization properties with neural networks. I hope the authors could provide some understandings of neural networks from the conclusion of the paper. Theoretical understanding is important, but it is better with a practical understanding.*

We will add a paragraph to the conclusions clarifying the general positioning of our work. Our setting is a shallow neural network in the feature learning regime - the same regime considered in the works on the mean-field limit. As argued in a number of works in this literature, we believe most of the state-of-the-art neural nets operate in such a feature rich regime. This is in stark contrast to the NTK regime, which ultimately reduces to a convex kernel method. Hence, NTK is considerably simpler from both the optimization and analysis point of view, and fails to capture feature learning, which is one of the key aspects of state-of-the-art neural networks.

> *Which techniques used in the paper do you think can be applied beyond the simple setting?*

We do believe that an analogous ODE description can be obtained in future work for deeper neural nets where the correlation structure between weights is more involved.

 Thank you for the answers.
I understand the practice in theoretical works.
Is there anything in common between properties of neural networks and the observations in the paper? There are many simple models, but why do you consider the settings in the paper. Is it just for analytic investigation? Does it have similar properties with neural networks? Is there anything similar to the convergence trade-off observed in neural networks? For example, the neural tangent kernel (NTK) has similar optimization properties with neural networks. I hope the authors could provide some understandings of neural networks from the conclusion of the paper. Theoretical understanding is important, but it is better with a practical understanding.

> the analysis of simple models constructively leads to the development of new techniques and concepts that will one day advance the field in a direction the practitioners did not anticipate. 

Which techniques used in the paper do you think can be applied beyond the simple setting?

As many other reviewers, I support accepting this paper. 

 We thank the reviewer for the constructive comments. Below we address the remarks.

> *1. In the introduction, the link between Saad \& Solla ODE and the mean-field PDE is mentioned (...). 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\ll pd/(d+\log (p))$ is enough for SGD to concentrate on the population gradient flow (perfect learning). This deserves (...).*

We do not know which works on the mean-field and the propagation-of-error the referee means here, could you please provide a pointer to this result?  
   
In our work we simply want to connect the regime where $d\to\infty$ and $p$ is fixed (Saad \& Solla) to the regime where $p\to\infty$ (mean-field). We thus extend the Saad \& Solla ODEs to the large $p$ regime.
   
The sentence in line 146 is indeed misleading (that paragraph discusses the relation of our proof to the proof in [6]) and we removed it.
   
We will comment on the NTK limit, which is very different lazy regime. We will also add a comment that the abc parametrization reduces to the mean-field scaling in the setting of two-layer networks (if we understood it correctly).

> *2. In 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).*

We thank the reviewer for pointing this out. Indeed, there is a very small term corresponding to the fluctuations $\| \| x \| \|^2/d$ in the update equations, that we chose to neglect. We will include it in the revision of the paper.
   
Fortunately, this term does not change the ODE approximation theorem: the conditions of Theorem A.1 (and especially the one regarding ${ \bf\Gamma}^\nu$) imply that a perturbation of order $O(d^{1/2})$ of Eqs. (11) does not affect the final bound of Theorem 2.1. A rigorous version of this justification will also be included in the revision.

> *3. 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.*

We agree with the reviewer interpretation. The "learning term" corresponds precisely to the gradient of the population risk. On the other hand, the "noise term" can be seen as a high-dimensional variance correction which obstructs learning in our setting. Thus our "learning" versus "noise" decomposition can indeed be interpreted as a "bias" versus "variance" decomposition. As discussed in the reply to reviewer *Q7z1*, it can be shown that asymptotically the "noise term" causes the population risk to plateau at $\Delta\gamma$ [6]. Therefore, in the noiseless regime the dynamics reaches zero risk. We will expand this discussion in the revised manuscript.

> *4. More 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$ ?*

This is a very good question. In our point of view, the existence of the ODEs has value, if only in making the plot in Fig. 1. We believe more along the lines of this question can be extracted and these are precisely the problems that we are focusing on in ongoing work.

> *5. Contrary to the discussion below Theorem 2.1, I do not think Assumption 3 is mild: for example, it cannot be true if $p>d$.*

Please see our answer to referee [JBZo](https://openreview.net/forum?id=GL-3WEdNRM&noteId=7EgpxQWmn9) who asked the same question.

> *6. Even 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?*

This is a very good question, that we indeed want to explore further. Recent works (in particular [27, 28]) show that when the gradient towards the teacher is non-vanishing (which corresponds to initial overlaps bounded away from zero), convergence happens in finite time. When this is not the case (e.g. uniform Gaussian initialization), and if the initial starting point of the population loss is a saddle, then a search phase of length $O (\log(d))$ happens first.

 We thank the reviewer for the constructive comments. Below we address the questions.

> *1. The authors use a scaling of the learning rate $\gamma$ like $\gamma = \gamma_0 / d^\delta $ with a constant  $\gamma_0 $ and a parameter $\delta $ that can be either positive or negative. A negative parameter 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.*

We agree that the negative exponent is unusual, this is why we found it rather intriguing that it could yield good results anyway.

As is well known, in practice, if one chooses a learning rate that is too large, the model can quickly converge to a suboptimal solution, whereas too small learning rate can make the training get stuck. Our results here indicates that one can leverage this trade-off by increasing the hidden layer, which presumably changes the optimisation landscape reducing the number of suboptimal solutions, allowing the increase of the learning rate.

> *2. 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.*

Comparing the SGD ran on real data and  with Gaussian ODEs has been done in previous papers, see e.g. the discussion lines 114-121. We invite for instance  the reviewer to consult Fig. 2 and Fig. 4 in Ref. [17].

 We thank the reviewer for the constructive comments. Below we address the remarks.

> *However, I have a concern about Assumption 3 made in Theorem 2.1 which assumes $\bar{{\bf \Omega}}$ 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
$\bar{{\bf \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, ${\bf W}^{\nu}$ will degenerate by the convergence to the teacher model.*

We indeed discuss the seemingly tricky assumption $3$ on the lines 147-153. Let us restate some of the comments here:

- Simulations always match extremely well with the ODE, so the positivity assumption is probably spurious.
- Given we work always with  $\Delta > 0$, there is always a slight, even tiny, mismatch between the vectors at finite time. Note that this is *also* the case in the mean-field limit, as discussed for instance in [15]: in the space orthogonal to the teacher function, the vectors remain far from each other during the iteration. 
- As we write in lines 147-153, one could technically get rid of this assumption by adding an infinitesimal noise to each set of weights at each step of the gradient iteration. This is -after all- what computers do beyond, say, $10^{-10}$ machine precision. 

We will add [Yang and Hu (2021)] to the related work discussion. Indeed also Referee *dAar* pointed it out as well.


 We thank the reviewer for the constructive comments. Below we address the numbered remarks.

> *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.*

Our contribution is to extend the Saad \& Solla-like approach to the wide neural networks. The implications of this on the works focusing on the mean-field regime is indeed a very interesting question. Such a comparison was in our point of view enabled by the work in ref [15] where the mean-field PDEs are reduced to a "dimension-free" form. With this form an explicit mapping between the ODEs and PDEs should be possible, but this is beyond the scope of the present paper. 

> *2. 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?*

The choice to fix the second layer weights to $1/p$ is indeed mostly used to simplify the analysis (and connect to the mean-field regime that uses the same convention). We do expect the same approximation results to hold in the general case where both layers are trained. To keep an almost dimension-free bound in $p$ in the ODE approximation, we would need to check that the second layer weights stay as $O(1)$ (in the mean-field scaling) throughout the trajectory, which seems plausible if they are initialized uniformly. See also reply to reviewer [Q7z1](https://openreview.net/forum?id=GL-3WEdNRM&noteId=HDLivAp2AGf) (item 5). We will add a discussion on these aspects.

 We thank the reviewer for the constructive comments. Below we address the remarks.

> *The only weakness I found is the setting the paper considered. I'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. It would be better for a broad range of community if the authors could provide further implications between the two settings:
(i) convex, two-layer, global convergence, with teacher-generated labels, limit of high-dimensional data, one-pass, Gaussian data (ii) non-convex, multi-layer, no global convergence, without teacher, finite high-dimensional data , (not one-pass?), general non-Gaussian data.*

As is the practice in theoretical works, one starts by looking at the minimal model capturing the complex phenomenology of interest which is not yet fully understood. In our work, we are interested in understanding the SGD dynamics in non-convex neural networks. The simplest model, which has attracted a lot of attention recently [1-16] but it is not yet fully understood, is the case of two-layer networks. This is indeed close to the setting (i) listed by the reviewer. However, we would like to stress that this learning problem is *not* convex, and this is precisely why it is an a challenge to understand the SGD dynamics for two-layer neural networks. 

Although we agree with the reviewer that going beyond two-layers and towards more "practical settings" is an important endeavour, the overall hope is that the analysis of simple models constructively leads to the development of new techniques and concepts that will one day advance the field in a direction the practitioners did not anticipate. We believe this view is widespread among other colleagues working in theoretical machine learning.

> *1. L17: Why is it a convex problem?*

As discussed above, the minimization problem for two-layer neural networks:

\begin{equation}
  \underset{ {\bf W} \in \mathbb{R}^{p\times d}}{\text{min}} \mathbb{E}
\left[ y -\frac{1}{p} \sum_j^{p} \sigma( {\bf w}_j^\top {\bf x})   \right]^2
\end{equation}
is *not* convex on the parameter/weight space ${\bf W}\in\mathbb{R}^{p\times d}$. The key observation in the works [1-4] cited in L17 is that , by defining the following empirical measure for the weights:

\begin{equation}
 \hat{\rho}_p ( {\bf w}) = \frac{1}{p}\sum_j^{p} \delta( {\bf w}-   {\bf w}_j )  .
\end{equation}
The minimization problem  can be re-written as:

\begin{equation}
  \underset{\hat{\rho}_p}{\text{min}} ~  \mathbb{E}  \left[ y -  \int \text{d} \hat{\rho}_p ( {\bf w } ) \sigma ( {\bf w }^\top  {\bf x }    )  \right]^2
\end{equation}
which is convex in the space of measures, since the integral is linear in $\hat{\rho}_p$. In [1-4], the authors exploit this fact to prove global convergence results in the limit $p\to\infty$. One of our goals in this work is investigate the intermediate regime where $p$ scales with $d$.

>*2. What does it exactly mean by "the one-pass SGD"?*

One-pass SGD means that the network is given a *single* sample $( {\bf x}^\nu  , y^\nu ) $ at a time and the weights are updated sequentially. The training examples are never revisited, i.e. $ \{ ( {\bf x}^\nu  , y^\nu ) \}_{\nu \ge 1} $ are i.i.d from $ \mathbb{P}  ( {\bf x} ,  y) $. This regime is studied in many theoretical works [1-16], and is not far from true for large scale applications, where the data set is so large that each sample is visited at visited at most a few times (see e.g. Bottou '10).

>*3. L202: What does it mean by 'a convergence trade-off'?*

We mean by the case where the faster one of contributions of Eqs. (16) goes to zero, the worse is the convergence rate from Theorem 2.1. For the bad learning regime  ($ - 1/2  < \kappa + \delta < 0$) , the learning term vanishes with $ d^{\kappa+\delta} $, while the convergence rate scales with $ d^{-(1/2+\kappa+\delta)} \log(d)$. Summarily, the more the learning term is attenuated, i.e. the more negative is $\kappa+\delta$, the worse the dynamics is described by the ODEs in this case. As a result, when we compare with simulations, this regime presents strong finite-size effects, as shown in Fig. 5 (left) . We have clarified this point in the revised version.

 We thank the reviewer for the constructive comments. Below we address the numbered remarks.

> *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.*

The function $\psi$ actually does not depend on $d$, except through the scaling of $p, \delta t$. We will comment on this dependency in the revised paper, thank you for the pointer.

> *2. Expand the literature on fluid-limit/ODE limits for analyzing randomized algorithms/Markov chain.*

We thank the reviewer for point us to this relevant literature and for the suggestion. The literature on fluid-limit/ODE limits will be expanded in the camera-ready version.

> *3. 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).*

We agree, and again thank the  reviewer for the suggestion. We will add this note in the post-theorem discussion.

> *4. 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.*

We completely agree with the reviewer that the discussion on the trade-off between "learning" and "noise" terms is crucial to understanding when perfect learning is possible. From a technical viewpoint, it can be shown that in the $\delta+\kappa=0$ regime the "noise terms" is responsible for the population risk to plateau at $\Delta\gamma$ [6] (hence the choice of naming). On a more intuitive level, the "learning term" is simply the gradient of the population risk, for which the global minimum is achieved at perfect learning, while the noise term is a high-dimensional variance correction which hinders learning. We have added a small note in the revised manuscript (due to the page limit) and will expand the discussion of this trade-off in the camera-ready version.

> *5. 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)]?*

The choice to fix the second layer weights to $1$ is indeed mostly used to simplify the analysis (it is also commonly done in the mean-field limit [1-4]), and we expect the same approximation results to hold in the general case where both layers are trained. Letting ${\bf a}^* \in  {\mathbb R}^k$ and $ {\bf a} \in {\mathbb R}^p$ denote the weights of the second layer (in the mean-field parametrization), the update equations would look like
\begin{equation}
   q_{jl}^{\nu+1} - q_{jl}^{\nu}  =  \frac{\gamma}{p d} \left(   a_j^{\nu}  {\cal E}_j^\nu \lambda_l^\nu + a_l^{\nu}  {\cal E}_l^\nu \lambda_j^\nu   \right) +  \frac{\gamma^2}{p^2 d}   a_j^{\nu} a_l^{\nu} {\cal E}_j^\nu  {\cal E}_l^\nu  ,
\end{equation}

\begin{equation}
   m_{jr}^{\nu+1} - m_{jr}^{\nu}= \frac{\gamma}{p d} a_j^{\nu} {\cal E}_j^\nu  \lambda_r^{*\nu}   ,
\end{equation}

\begin{equation}
   a_j^{\nu+1} -  a_j^\nu= \frac{\tilde{\gamma}}{p} {\cal E}^\nu   ,
\end{equation}
with $ {\cal E}_j^\nu \equiv \sigma'( \lambda_j^\nu )  {\cal E}^\nu $ and now

\begin{equation}
     {\cal E}^\nu\equiv \frac{1}{k} \sum_{r=1}^{k} a_r^{*\nu}  \sigma( \lambda_{r}^{*\nu}  ) - \frac{1}{p} \sum_{l=1}^{p}  a_l^\nu \sigma ( \lambda_{l}^\nu )   + \sqrt{\Delta} \zeta^\nu .
\end{equation}

To keep an almost dimension-free bound in $p$ in the ODE approximation, we would need to check that the second layer weights stay as $ {\cal O} (1)$ throughout the trajectory, which seems plausible if they are initialized uniformly.

 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*.
They 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.
The 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.

 Strengths:

- Scaling of learning rate and number of hidden units in terms of the dimension is a real improvement over the Saad and Sola setting.
- the paper is quite well written.
- So far I checked, the proofs are ok.

Weaknesses:

- 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. 
- 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. - 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.

- 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. Not applicable.

 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.
 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.  I list below a few comments, questions and concerns.

In 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.

In 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.

More 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$ ? 

Contrary to the discussion below Theorem 2.1, I do not think Assumption 3 is mild: for example, it cannot be true if $p > d$.  

Even 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?

Line 206: *leads Yes.

 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.  **Strengths** 

It 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. 

The 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. 

**Weaknesses**

However, 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.


**Other Comments**
The 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. 

[Yang and Hu (2021)] Greg Yang and Edward J. Hu. Feature Learning in Infinite-Width Neural Networks. ICML, 2021. 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. 
If my understanding is wrong, I would appreciate it if you could correct me.
 N/A

 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.  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. 

Quality and Clarity: The paper is very well written. All the mathematics and discussions are clear and easy to read.

Significance: 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.  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. 

2. 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? The authors have honestly discussed the limitations of the work. 

 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$. ### Strength
- very rigorous approach
- much more generalized settings compared to previous works (e.g., Saad & Solla 1995).

### Weaknesses
- The only weakness I found is the setting the paper considered. 
I'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.
It would be better for a broad range of community if the authors could provide further implications between the two settings:
- (i) convex, two-layer, global convergence, with teacher-generated labels, limit of high-dimensional data $d\rightarrow \infty$, one-pass, Gaussian data
- (ii) non-convex, multi-layer, no global convergence, without teacher, finite high-dimensional data $d<\infty$, (not one-pass?), general non-Gaussian data
 - L17: Why is it a convex problem?
- What does it exactly mean by "the one-pass SGD"?
- L202: What does it mean by "a convergence trade-off"?
 Please provide further limitations when introducing the specific setting considered in the paper.

 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. 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$.

As weaknesses, I have only a few relatively minor remarks that are in next section. 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.

2. 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.

3. 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).

4. 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.

5. 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)]?

**[responses by the authors acknowledged (4/8/2022)]**
Thank you for these clarifications.

 NA.

Meta-review: The reviewers and I agree that the contributions of the paper are of interest and useful addition to the literature. Therefore, I recommend accepting the paper.

Please consider the reviewers' comments when preparing the camera-ready version.


Decision: Accept