Learning time-scales in two-layers neural networks
Gradient-based learning in multi-layer neural networks displays a number of
striking features. In particular, the decrease rate of empirical risk is
non-monotone even after averaging over large batches. Long plateaus in which
one observes barely any progress alternate with intervals of rapid decrease.
These successive phases of learning often take place on very different time
scales. Finally, models learnt in an early phase are typically `simpler' or
`easier to learn' although in a way that is difficult to formalize.
  Although theoretical explanations of these phenomena have been put forward,
each of them captures at best certain specific regimes. In this paper, we study
the gradient flow dynamics of a wide two-layer neural network in
high-dimension, when data are distributed according to a single-index model
(i.e., the target function depends on a one-dimensional projection of the
covariates). Based on a mixture of new rigorous results, non-rigorous
mathematical derivations, and numerical simulations, we propose a scenario for
the learning dynamics in this setting. In particular, the proposed evolution
exhibits separation of timescales and intermittency. These behaviors arise
naturally because the population gradient flow can be recast as a singularly
perturbed dynamical system.
Raphaël Berthier, Andrea Montanari, Kangjie Zhou
# Learning time-scales in two-layers neural networks

###### Abstract

Gradient-based learning in multi-layer neural networks displays a number of striking features. In particular, the decrease rate of empirical risk is non-monotone even after averaging over large batches. Long plateaus in which one observes barely any progress alternate with intervals of rapid decrease. These successive phases of learning often take place on very different time scales. Finally, models learnt in an early phase are typically'simpler' or 'easier to learn' although in a way that is difficult to formalize.

Although theoretical explanations of these phenomena have been put forward, each of them captures at best certain specific regimes. In this paper, we study the gradient flow dynamics of a wide two-layer neural network in high-dimension, when data are distributed according to a single-index model (i.e., the target function depends on a one-dimensional projection of the covariates). Based on a mixture of new rigorous results, non-rigorous mathematical derivations, and numerical simulations, we propose a scenario for the learning dynamics in this setting. In particular, the proposed evolution exhibits separation of timescales and intermittency. These behaviors arise naturally because the population gradient flow can be recast as a singularly perturbed dynamical system.

###### Contents

* 1 Introduction
* 2 Setting and standard learning scenario
* 3 Further related work
* 4 The large-network, high-dimensional limit
	* 4.1 Connection with mean field theory
	* 4.2 A general formulation
* 5 Numerical solution
* 6 Timescales hierarchy in the gradient flow dynamics
	* 6.1 First time scale: constant component
	* 6.2 Second time scale: linear component I
	* 6.3 Third time scale: linear component II
	* 6.4 Conjectured behavior for larger time scales
* 7 Stochastic gradient descent and finite sample size
A Appendix to Section 4A1 Proof of Proposition 1A.2 Proof of Corollary 1A.3 Proof of Proposition 2A.4 Derivation of the mean field dynamics (28)A.5 Details of the alternative mean field approachB Calculations for the analysis of mean-field gradient flowB.1 Solution of Eq. (83)B.2 Induced approximation of the riskB.3 Proof of Theorem 1C Proofs of Theorem 2 and 3: learning with projected SGDC.1 Difference between GF and GDC.2 Difference between GD and SGDC.3 Difference between SGD and projected SGDC.4 Proof of Theorem 3D Counterexamples to the standard learning scenarioD.1 Case 1: \(\sigma_{k}=0\) for some \(k\in\mathbb{N}\)D.2 Case 2: \(\varphi_{0}=\dots=\varphi_{k}=0\) for some \(k\geq 1\)D.3 Case 3: \(\varphi_{k}=0\) for some \(k\geq 1\)

## 1 Introduction

It is a recurring empirical observation that the training dynamics of neural networks exhibits a whole range of surprising behaviors:

1. _Plateaus._ Plotting the training and test error as a function of SGD steps, using either small stepsize or large batches to average out stochasticity, reveals striking patterns. These error curves display long plateaus where barely anything seems to be happening, which are followed by rapid drops (Saad and Solla, 1995; Yoshida and Okada, 2019; Power et al., 2022).
2. _Time-scales separation._ The time window for this rapid descent is much shorter than the time spent in the plateaus. Additionally, subsequent phases of learning take increasingly longer times (Ghorbani et al., 2020; Barak et al., 2022).
3. _Incremental learning._ Models learnt in the first phases of learning appear to be simpler than in later phases. Among others, Arpit et al. (2017) demonstrated that easier examples in a dataset are learned earlier; Kalimeris et al. (2019) showed that models learnt in the first phase of training correlate well with linear models; Gissin et al. (2019) showed that, in many simplified models, the dynamics of gradient descent explores the solution space in an incremental order of complexity; Power et al. (2022) demonstrated that, in certain settings, a function that approximates well the target is only learnt past the point of overfitting.

Understanding these phenomena is not a matter of intellectual curiosity. In particular, incremental learning plays a key role in our understanding of generalization in deep learning. Indeed, in this scenario, stopping the learning at a certain time \(t\) amounts to controlling the complexity of the model learnt. The notion of complexity corresponds to the order in which the space of models is explored.

While a number of groups have developed models to explain these phenomena, it is fair to say that a complete picture is still lacking. An exhaustive overview of these works is out of place
here. We will outline three possible explanations that have been developed in the past, and provide more pointers in Section 3.

Theory \(\#1\): Dynamics near singular points.Several early works (Saad and Solla, 1995; Fukumizu and Amari, 2000; Wei et al., 2008) pointed out that the parametrization of multi-layer neural networks presents symmetries and degeneracies. For instance, the function represented by a multilayer perceptron is invariant under permutations of the neurons in the same layer. As a consequence, the population risk has multiple local minima connected through saddles or other singular sub-manifolds. Dynamics near these sub-manifolds naturally exhibits plateaus. Further, random or agnostic initializations typically place the network close to such submanifolds.

Theory \(\#2\): Linear networks.Following the pioneering work of Baldi and Hornik (1989), a number of authors, most notably Saxe et al. (2013), Li et al. (2020), studied the behavior of deep neural networks with linear activations. While such networks can only represent linear functions, the training dynamics is highly non-linear. As demonstrated in Saxe et al. (2013), learning happens through stages that correspond to the singular value decomposition of the input-output covariance. Time scales are determined by the singular values.

Theory \(\#3\): Kernel regime.Following an initial insight of Jacot et al. (2018), a number of groups proved that, for certain initializations, the training dynamics and model learnt by overparametrized neural networks is well approximated by certain linearly parametrized models. In the limit of very wide networks, the training dynamics of these models converges in turn to the training dynamics of kernel ridge(less) regression (KRR) with respect to a deterministic kernel (independent of the random initialization.) We refer to Bartlett et al. (2021) for an overview and pointers to this literature. Recently Ghosh et al. (2021) show that, in high dimension, the learning dynamics of KRR also exhibits plateaus and waterfalls, and learns functions of increasing complexity over a diverging sequence of timescales.

While each of these theories offers useful insights, it is important to realize that they do not agree on the basic mechanism that explains plateaus, time-scales separation, and incremental learning. In theory \(\#1\), plateaus are associated to singular manifolds and high-dimensional saddles, while in theories \(\#2\) and \(\#3\) they are related to a hierarchy of singular values of a certain matrix. In \(\#2\), the relevant singular values are the ones of the input-output covariance, and the fact that these singular values are well separated is postulated to be a property of the data distribution. In contrast, in \(\#3\) the relevant singular values are the eigenvalues of the kernel operator, and hence completely independent of the output (the target function). In this case, eigenvalues which are very different are proved to exist under natural high-dimensional distributions.

Not only these theories propose different explanations, but they are also motivated by very different simplified models. Theory \(\#1\) has been developed only for networks with a small number of hidden units. Theory \(\#2\) only applies to networks with multiple output units, because otherwise the input-output covariance is a \(d\times 1\) matrix and hence has only one non-trivial singular value. Finally, theory \(\#3\) applies under the conditions of the linear (a.k.a. lazy) regime, namely large overparametrization and suitable initialization (see, e.g., Bartlett et al. (2021)).

In order to better understand the origin of plateaus, time-scales separation, and incremental learning, we attempt a detailed analysis of gradient flow for two-layer neural networks. We consider a simple data-generation model, and propose a precise scenario for the behavior of learning dynamics. We do not assume any of the simplifying features of the theories described above: activations are non-linear; the number of hidden neurons is large; we place ourselves outside the linear (lazy) regime.
Our analysis is based on methods from dynamical systems theory: singular perturbation theory and matched asymptotic expansions. Unfortunately, we fall short of providing a general rigorous proof of the proposed scenario, but we can nevertheless prove it in several special cases and provide a heuristic argument supporting its generality.

The rest of the paper is organized as follows. Section 2 describes our data distribution, learning model, and the proposed scenario for the learning dynamics. We review further related work in Section 3. Section 4 describes the reduction of the gradient flow to a'mean field' dynamics that will be the starting point of our analysis. Section 5 presents numerical evidence of the proposed learning scenario. Finally, Sections 6 to 7 present our analysis of the learning dynamics.

Notations.In this paper, we use the classical asymptotic notations. The notations \(f(\varepsilon)=o(g(\varepsilon))\) or \(g(\varepsilon)=\omega(f(\varepsilon))\) as \(\varepsilon\to 0\) both denote that \(|f(\varepsilon)|/|g(\varepsilon)|\to 0\) in the limit \(\varepsilon\to 0\). The notations \(f(\varepsilon)=O(g(\varepsilon))\) or \(g(\varepsilon)=\Omega(f(\varepsilon))\) both denote that the ratio \(|f(\varepsilon)|/|g(\varepsilon)|\) remains upper bounded in the limit. The notation \(f(\varepsilon)=\Theta(g(\varepsilon))\) or \(f(\varepsilon)\asymp g(\varepsilon)\) denote that \(f(\varepsilon)=O(g(\varepsilon))\) and \(g(\varepsilon)=O(f(\varepsilon))\) both hold. Finally, \(f(\varepsilon)\sim g(\varepsilon)\) denotes that \(f(\varepsilon)/g(\varepsilon)\to 1\) in the limit.

## 2 Setting and standard learning scenario

We are given pairs \(\{(x_{i},y_{i})\}_{i\leq n}\), where \(x_{i}\in\mathbb{R}^{d}\) is a feature vector and \(y_{i}\in\mathbb{R}\) is a response variable. We are interested in cases in which the feature vector is high-dimensional but does not contain strong structure, but the response depends on a low-dimensional projection of the data. We assume the simplest model of this type, the so-called single-index model:

\[y_{i}=\varphi(\langle u_{*},x_{i}\rangle)\,,\qquad\,x_{i}\sim \mathsf{N}(0,I_{d}),\;u_{*}\in\mathbb{S}^{d-1}, \tag{1}\]

where \(\varphi:\mathbb{R}\to\mathbb{R}\) is a link function, \(\mathsf{N}(0,I_{d})\) denotes the standard multivariate Gaussian distribution in dimension \(d\), and \(\mathbb{S}^{d-1}:=\{v\in\mathbb{R}^{d}:\,\|v\|_{2}=1\}\). We study the ability to learn model (1) using a two-layers neural network with \(m\) hidden neurons:

\[f(x;a,u)=\frac{1}{m}\sum_{i=1}^{m}a_{i}\sigma(\langle u_{i},x\rangle),\qquad a _{1},\cdots,a_{m}\in\mathbb{R},\;u_{1},\cdots,u_{m}\in\mathbb{S}^{d-1},\]

where \((a,u):=(a_{1},\cdots,a_{m},u_{1},\cdots,u_{m})\) collectively denotes all the model's parameter. The factor \(1/m\) in the definition is relevant for the initialization and learning rate. We anticipate that we will initialize the \(a_{i}\)'s to be of order one, which results in second layer coefficients \(a_{i}/m=\Theta(1/m)\). This is often referred to as the'mean-field initialization' and is known to drive learning process out of the linear or kernel regime, see e.g. (Mei et al., 2018; Chizat and Bach, 2018; Ghorbani et al., 2020; Yang and Hu, 2020; Abbe et al., 2022).

The bulk of our work will be devoted to the analysis of projected gradient flow in \((a_{i},u_{i})_{1\leq i\leq m}\) on the population risk

\[\mathscr{R}(a,u) =\frac{1}{2}\mathbb{E}\big{\{}\big{(}y-f(x;a,u)\big{)}^{2}\big{\}} \tag{2}\] \[=\frac{1}{2}\mathbb{E}\Big{\{}\Big{(}\varphi(\langle u_{*},x \rangle)-\frac{1}{m}\sum_{i=1}^{m}a_{i}\sigma(\langle u_{i},x\rangle)\Big{)}^{ 2}\Big{\}}\,. \tag{3}\]

In Section 7, we will bound the distance between stochastic gradient descent (SGD) and gradient flow in population risk. As a consequence, we will establish finite sample generalization guarantees for SGD learning.

Projected gradient flow with respect to the risk \(\mathscr{R}(a,u)\) is defined by the following ordinary differential equations (ODEs):

\[\partial_{t}(\varepsilon a_{i})=-m\partial_{a_{i}}\mathscr{R}(a,u)\,, \tag{4}\]
\[\partial_{t}u_{i}=-m(I_{d}-u_{i}u_{i}^{\top})\nabla_{u_{i}}\mathscr{R}(a,u)\,. \tag{5}\]

It is useful to make a few remarks about the definition of gradient flow:

* The projection \(I_{d}-u_{i}u_{i}^{\top}\) ensures that \(u_{i}\) remains on the unit sphere \(\mathbb{S}^{d-1}\).
* The overall scaling of time is arbitrary, and the matching to SGD steps will be carried out in Section 7. The factors \(m\) on the right-hand side are introduced for mathematical convenience, since the partial derivatives are of order \(1/m\).
* The factor \(\varepsilon\) introduced in the flow of the \(a_{i}\)'s reflects the fact that usually SGD is run with respect to the overall second-layer weights \((a_{i}/m)_{1\leq i\leq m}\). This would correspond to taking \(\varepsilon=1/m\). However, we will keep \(\varepsilon\) as a free parameter independent of \(m\), and study the evolution for small \(\varepsilon\).

We assume the initialization to be random with i.i.d. components \((a_{i,\mathrm{init}},u_{i,\mathrm{init}})\):

\[(a_{i,\mathrm{init}},u_{i,\mathrm{init}})\sim\mathrm{P}_{A}\otimes\mathrm{ Unif}(\mathbb{S}^{d-1})\,, \tag{6}\]

where \(\mathrm{P}_{A}\) is a probability measure on \(\mathbb{R}\). The unique solution of the gradient flow ODEs with this initialization will be denoted by \((a(t),u(t))\). We will be interested in the case of large networks \((m\to\infty)\) in high dimension \((d\to\infty)\). As shown below, the two limits commute (over fixed time horizons).

Our main finding is that, in a number of cases, \(\varphi\) is learnt incrementally. Namely, the function \(f(x;a(t),u(t))\) evolves over time according to a sequence of polynomial approximations of \(\varphi((u_{*},x))\). These polynomial approximations are given by the decomposition of \(\varphi\) in \(L^{2}(\mathbb{R},\phi(x)\mathrm{d}x)\), where \(\phi(x)\) is the standard normal density: \(\phi(x)=\exp(-x^{2}/2)/\sqrt{2\pi}\). (For notational simplicity, we will use the shorthand \(L^{2}\) instead of \(L^{2}(\mathbb{R},\phi(x)\mathrm{d}x)\) in the sequel.)

In order to describe the polynomial approximations learnt during the training more explicitly, we decompose \(\varphi\) and \(\sigma\) into normalized Hermite polynomials:

\[\varphi(z)=\sum_{k=0}^{\infty}\varphi_{k}\mathrm{He}_{k}(z)\,,\quad\sigma(z)= \sum_{k=0}^{\infty}\sigma_{k}\mathrm{He}_{k}(z)\,. \tag{7}\]

Here, \(\mathrm{He}_{k}\) denotes the \(k\)-th Hermite polynomial, normalized so that \(\|\mathrm{He}_{k}\|_{L^{2}(\mathbb{R},\phi(x)\mathrm{d}x)}=1\).

As we will see, the incremental learning behavior arises for small \(\varepsilon\). By the law of large numbers (see below), the following almost sure limit exists (provided \(\mathrm{P}_{A}\) is square integrable)

\[\mathscr{R}_{\mathrm{init}}:=\lim_{m\to\infty}\lim_{d\to\infty}\mathscr{R}(a _{\mathrm{init}},u_{\mathrm{init}})\,=\frac{1}{2}\left(\varphi_{0}-\sigma_{0} \int\!a\,\mathrm{P}_{A}(\mathrm{d}a)\right)^{2}+\frac{1}{2}\sum_{k\geq 1} \varphi_{k}^{2}. \tag{8}\]

We are now in position to describe the scenario that we will study in the rest of the paper.

**Definition 1**.: _We say that the standard learning scenario holds up to level \(L\) for a certain target function \(\varphi\), activation \(\sigma\), and distribution \(\mathrm{P}_{A}\), if the followings hold:_

1. _The limit below exists:_ \[\mathscr{R}_{\infty}(t,\varepsilon)=\lim_{m\to\infty}\lim_{d\to\infty}\mathscr{ R}(a(t),u(t)).\] (9)
2. _There exist constants_ \(c_{2},\ldots,c_{L+1}>0\) _such that the following asymptotic holds as_ \(\varepsilon\to 0\)_,_ \(t\to 0\)_:_ \[\mathscr{R}_{\infty}(t,\varepsilon)\xrightarrow[\varepsilon\to 0,t\to 0]{}\begin{cases} \mathscr{R}_{\mathrm{init}}&\text{if }t=o(\varepsilon)\,,\\ \frac{1}{2}\sum_{k\geq 1}\varphi_{k}^{2}&\text{if }t=\omega(\varepsilon)\text{ and }t=\frac{1}{4|\sigma_{1} \varphi_{1}|}\varepsilon^{\nicefrac{{1}}{{2}}}\log\frac{1}{\varepsilon}- \omega(\varepsilon^{\nicefrac{{1}}{{2}}})\,,\\ \frac{1}{2}\sum_{k\geq 2}\varphi_{k}^{2}&\text{if }t=\frac{1}{4|\sigma_{1} \varphi_{1}|}\varepsilon^{\nicefrac{{1}}{{2}}}\log\frac{1}{\varepsilon}+ \omega(\varepsilon^{\nicefrac{{1}}{{2}}})\text{ and }t=c_{2}\varepsilon^{\nicefrac{{1}}{{4}}}- \omega(\varepsilon^{\nicefrac{{1}}{{3}}})\,,\\ \frac{1}{2}\sum_{k\geq l}\varphi_{k}^{2}&\text{if }t=c_{l-1}\varepsilon^{ \nicefrac{{1}}{{2(l-1)}}}+\omega(\varepsilon^{\nicefrac{{1}}{{l}}})\text{ and }t=c_{l}\varepsilon^{\nicefrac{{1}}{{2l}}}- \omega(\varepsilon^{\nicefrac{{1}}{{l+1}}})\,,\\ &\text{ for all }3\leq l\leq L+1.\end{cases}\]
Figure 1 provides a cartoon illustration of the standard learning scenario.

A specific realization of our general setup is determined by the triple \((\sigma,\varphi,\mathrm{P}_{A})\), In the rest of the paper, we will provide evidence showing that the standard learning scenario holds in a number of cases. Nevertheless, we can also construct examples in which it does not hold:

* If one or more of the Hermite coefficients of the activation vanish, then the standard scenario does not hold for general \(\varphi\). Specifically, if \(\sigma_{k}=0\), then for any \(t\) the function \(f(x;a(t),u(t))\) remains orthogonal to \(\mathrm{He}_{k}(\langle u_{*},x\rangle)\). In particular, if \(\varphi_{k}\neq 0\) then the risk remains bounded away from zero for every \(t\). We refer to Appendix D.1 for a formal statement.
* If the first Hermite \(k+1\) coefficients of \(\varphi\) vanish, \(\varphi_{0}=\cdots=\varphi_{k}=0\), \(k\geq 1\), then the standard scenario does not hold. (See Appendix D.2 for the proof.)
* In fact, we expect the standard scenario might fail every time one or more of the coefficients \(\varphi_{k}\) vanish, for \(k\geq 1\). Appendix D.3 provides some heuristic justification for this failure.

**Remark 2.1**.: We can compare the standard learning scenario described here to the ones in earlier literature and described as theory #1, #2, #3 in the introduction. There appears points of contact, but also important differences with both theory #1 and #3:

* As in theory #1, the plateaus and separation of time scales arise because the trajectory of gradient flow is approximated by a sequence of motions along submanifolds in the space of parameters \((a,u)\). Along the \(l\)-th such submanifold \(f(x;a,u)\) is well-approximated by a degree-\(l\) polynomial. Escaping each submanifold takes an increasingly longer time. This is reminding of the motion between saddles investigated in earlier work (Saad and Solla, 1995; Fukumizu and Amari, 2000; Wei et al., 2008). However, unlike in earlier work, we will see that this applies to networks with a large (possibly diverging) number of

Figure 1: Cartoon illustration of the evolution of the population risk within the standard learning scenario of Definition 1.
hidden neurons. Also, we identify the subsequent phases of learning with the polynomial decomposition of Eq. (7).
* As in theory #3, subsequent phases of learning correspond to increasingly accurate polynomial approximations of the target function \(\varphi(\langle u_{*},x\rangle)\). However, the underlying mechanism and time scales are completely different. In the linear regime, the different time scales emerge because of increasingly small eigenvalues of the neural tangent kernel. In that case, the time required to learn degree-\(l\) polynomials is of order \(d^{l}\)(Ghosh et al., 2021). In contrast, in the standard learning scenario, polynomials of degree \(l\) are learnt on a time scale of order one in \(d\) (and only depending on the learning rate \(\varepsilon\)). This of course has important implications when approximating gradient flow by SGD. Within the linear regime, the sample size required to learn polynomial of order \(l\) scales like \(d^{l}\)(Ghosh et al., 2021), while in the standard scenario, it is only of order \(d\) (see Section 7).

## 3 Further related work

As we mentioned in the introduction, plateaus and time scales in the learning dynamics of kernel models were analyzed by Ghosh et al. (2021). A sharp analysis for the related random features model was developed by Bodin and Macris (2021).

Our analysis builds upon the mean-field description of learning in two-layer neural networks, which was developed in a sequence of works, see, e.g., (Mei et al., 2018; Rotskoff and Vanden-Eijnden, 2018; Chizat and Bach, 2018; Mei et al., 2019). In particular, we leverage the fact that, for the data distribution (1), the population risk function is invariant under rotations around the axis \(u_{*}\), and this allows for a dimensionality reduction in the mean field description. Similar symmetry argument were used by Mei et al. (2018) and, more recently, by Abbe et al. (2022).

The single-index model can be learnt using simpler methods than large two-layer networks. Limiting ourselves to the case of gradient descent algorithms, Mei et al. (2018) proved that gradient descent with respect to the non-convex empirical risk \(\widehat{R}_{n}(u):=n^{-1}\sum_{i=1}^{n}(y_{i}-\varphi(u^{\top}x_{i}))^{2}\) converges to a near global optimum, provided \(\varphi\) is strictly increasing. Ben Arous et al. (2021) considered online SGD under more challenging learning scenarios and characterized the time (sample size) for \(|\langle u,u_{*}\rangle|\) to become significantly larger than for a random unit vector \(u\).

Learning in overparametrized two-layer networks under model (1) (or its variations) has been studied recently by several groups. In particular, Ba et al. (2022) considers a training procedure which runs a single step gradient descent followed by freezing the first layer and performing ridge regression with respect to the second layer. This scheme is amenable to a precise characterization of the generalization error. Bietti et al. (2022) consider a similar scheme in which a first phase of gradient descent is run to achieve positive correlation with the unknown direction \(u_{*}\). Damian et al. (2022) also consider a two-phases scheme, and prove consistency and excess risk bounds for a more general class of target functions whereby the first equation in (1) is replaced by

\[y_{i}=\varphi(U_{*}^{\top}x_{i})+\varepsilon_{i}\,,\;\;\;U_{*}\in\mathbb{R}^{k \times d}\,,\varphi:\mathbb{R}^{k}\to\mathbb{R}\,, \tag{10}\]

with \(k\ll d\). In particular, near optimal error bounds are obtained under a non-degeneracy condition on \(\nabla^{2}\varphi\).

Abbe et al. (2022) consider a similar model whereby \(x\sim\text{Unif}(\{+1,-1\}^{d})\), and \(y=\varphi(x_{S})\) where \(S\subseteq[d]\), and \(x_{S}=(x_{i})_{i\in S}\) (i.e., \(x_{S}\) contains the coordinates of \(x\) indexed by entries of \(S\)). Under a structural assumption on \(\varphi\) (the'merged staircase property'), and for \(|S|\) fixed, they prove the two stages algorithm learns the target function with sample complexity of order \(d\). This paper is technically related to ours in that it uses mean-field theory to obtain a characterization of learning in terms of a PDE in a reduced \((k+2)\)-dimensional space.
A similar model was studied by Barak et al. (2022) that bounds the sample complexity by \(d^{O(k)}\) for learning parities on \(k\) bits using gradient descent with large batches (if \(k=O(1)\), Barak et al. (2022) require \(O(1)\) steps with batch size \(d^{O(k)}\)).

Let us emphasize that our objective is quite different from these works. We do not allow ourselves deviations from standard SGD and try to derive a precise picture of the successive phases of learning (in particular, we do not consider two-stage schemes or layer-by-layer learning). On the other hand, we focus on a relatively simple model.

To clarify the difference, it is perhaps useful to rephrase our claims in terms of sample complexity. While previous works show that the target function can be learnt with \(O(d)\) samples, we claim that it is learnt by online SGD with test error \(r\) from about \(C(r,\varepsilon)d\) samples and characterize the dependence of \(C(r,\varepsilon)\) on \(r\) for small \(\varepsilon\). (Falling short of a proof in the general case.)

After posting an initial version of this paper, we became aware that Arnaboldi et al. (2023) independently derived equations similar to (14)-(18), (25), (119). There are technical differences, and hence we cannot apply their results directly. However, Section 4.1 and Appendix A.4 are analogous to their work.

## 4 The large-network, high-dimensional limit

The first step of our analysis is a reduction of the system of ODEs (4), (5), with dimension \(m(d+1)\) to a system of ODEs in \(2m\) dimensions. We will achieve this reduction in two steps:

* First we reduce to a system in \(m(m+3)/2\) dimensions for the variables \(a_{i}\), \(\langle u_{i},u_{j}\rangle\), \(\langle u_{i},u_{*}\rangle\). This reduction is exact and is quite standard.
* We then show that the products \(\langle u_{i},u_{j}\rangle\) can be eliminated, with an error \(O(1/m)\). As further discussed below, the resulting dynamics could also be derived from the mean field theory of Mei et al. (2018); Rotskoff and Vanden-Eijnden (2018); Chizat and Bach (2018); Mei et al. (2019) (with the required modifications for the constraints \(\left\lVert u_{i}\right\rVert=1\)).

In order to define formally the reduced system, we define the functions \(U,V:[-1,1]\to\mathbb{R}\) via:

\[V(s) :=\mathbb{E}\{\varphi(G)\,\sigma(G_{s})\}=\sum_{k\geqslant 0} \varphi_{k}\sigma_{k}s^{k}\,,\hskip 14.226378pt(G,G_{s})\sim\mathcal{N}\left(0, \begin{bmatrix}1&s\\ s&1\end{bmatrix}\right)\,, \tag{11}\] \[U(s) :=\mathbb{E}\{\sigma(G)\,\sigma(G_{s})\}=\sum_{k\geqslant 0} \sigma_{k}^{2}s^{k}\,. \tag{12}\]

Note that the above identities follow from (O'Donnell, 2014, Proposition 11.31). Throughout this section, we will make the following assumptions.

* The distribution of weights at initialization, \(\mathrm{P}_{A}\) is supported on \([-M_{1},M_{1}]\).
* The activation function is bounded: \(\left\lVert\sigma\right\rVert_{\infty}\leq M_{2}\). Additionally, the functions \(V\) and \(U\) are bounded and of class \(C^{2}\), with uniformly bounded first and second derivatives over \(s\in[-1,1]\). A sufficient condition for this is \[\sup\left\{\left\lVert\sigma^{\prime}\right\rVert_{L^{2}},\,\left\lVert\sigma ^{\prime\prime}\right\rVert_{L^{2}}\right\}\leq M_{2},\hskip 14.226378pt\sup \left\{\left\lVert\varphi\right\rVert_{L^{2}},\,\left\lVert\varphi^{\prime} \right\rVert_{L^{2}},\,\left\lVert\varphi^{\prime\prime}\right\rVert_{L^{2}} \right\}\leq M_{2}.\]
* Responses are bounded, i.e., \(\left\lVert\varphi\right\rVert_{\infty}\leq M_{3}\).

**Remark 4.1**.: We hereby briefly explain the sufficiency of \(L^{2}\)-boundedness of derivatives of \(\sigma\) and \(\varphi\) as claimed in Assumption A2. Suppose for example that \(\left\lVert\sigma^{\prime}\right\rVert_{L^{2}}\leq M_{2}\) and \(\left\lVert\varphi^{\prime}\right\rVert_{L^{2}}\leq M_{2}\), then we have

\[\sup_{s\in[-1,1]}\left\lvert V^{\prime}(s)\right\rvert\overset{(a)}{=}\sup_{ s\in[-1,1]}\left\lvert\mathbb{E}\{\varphi^{\prime}(G)\,\sigma^{\prime}(G_{s})\} \right\rvert\overset{(b)}{\leq}\left\lVert\varphi^{\prime}\right\rVert_{L^{2} }\left\lVert\sigma^{\prime}\right\rVert_{L^{2}}\leq M_{2}^{2}, \tag{13}\]
where \((a)\) follows from Gaussian integration by parts and \((b)\) follows from Cauchy-Schwarz inequality.

Our first statement establishes reduction \((i)\) mentioned above. The proof of this fact is presented in Appendix A.1.

**Proposition 1** (Reduction to \(d\)-independent flow).: _Define \(s_{i}=\langle u_{i},u_{*}\rangle\), \(r_{ij}=\langle u_{i},u_{j}\rangle\) for \(i,j=1,\ldots,m\). Then, letting \(R=(r_{ij})_{i,j\leq m}\), we have_

\[\mathscr{R}(a,u)=\mathscr{R}_{\mathrm{red}}(a,s,R):=\frac{1}{2}\|\varphi\|_{L ^{2}}^{2}-\frac{1}{m}\sum_{i=1}^{m}a_{i}V(s_{i})+\frac{1}{2m^{2}}\sum_{i,j=1}^ {m}a_{i}a_{j}U(r_{ij})\,. \tag{14}\]

_If \((a(t),u(t))\) solve the gradient flow ODEs (4)-(5) then \((a(t),s(t),R(t))\) are the unique solution of the following set of ODEs (note that \(r_{ii}=1\) identically)_

\[\varepsilon\partial_{t}a_{i}= \,V(s_{i})-\frac{1}{m}\sum_{j=1}^{m}a_{j}U(r_{ij})\,, \tag{15}\] \[\partial_{t}s_{i}= \,a_{i}\left(V^{\prime}(s_{i})(1-s_{i}^{2})-\frac{1}{m}\sum_{j=1} ^{m}a_{j}U^{\prime}(r_{ij})(s_{j}-r_{ij}s_{i})\right)\,,\] (16) \[\partial_{t}r_{ij}= \,a_{i}\left(V^{\prime}(s_{i})(s_{j}-s_{i}r_{ij})-\frac{1}{m} \sum_{p=1}^{m}a_{p}U^{\prime}(r_{ip})(r_{jp}-r_{ip}r_{ij})\right)\,,\] (17) \[+a_{j}\left(V^{\prime}(s_{j})(s_{i}-s_{j}r_{ij})-\frac{1}{m}\sum_ {p=1}^{m}a_{p}U^{\prime}(r_{jp})(r_{ip}-r_{jp}r_{ij})\right)\,. \tag{18}\]

The input dimension \(d\) does not appear in the reduced ODEs, Eqs. (15) to (18), and only plays a role in the initialization of the \(s_{i}\)'s and the \(r_{ij}\)'s. Namely, since \(u_{i,\mathrm{init}}\sim\mathrm{Unif}(\mathbb{S}^{d-1})\), we can represent \(u_{i,\mathrm{init}}=g_{i}/\|g_{i}\|_{2}\) with \(g_{i}\sim\mathsf{N}(0,I_{d}/d)\). By concentration of \(\|g_{i}\|_{2}\), this implies that, for \(1\leq i<j\leq m\), \(s_{i}\), \(r_{ij}\) are approximately \(\mathsf{N}(0,1/d)\).

This discussion immediately yields the following consequence.

**Corollary 1**.: _Let \((a(t),u(t))\) be the solution of the gradient flow ODEs (4), (5) with initialization (6), and let \((a^{0}(t),s^{0}(t),R^{0}(t))\) be the unique solution of Eqs. (15) to (18), with initialization \(a^{0}_{i}(0)=a_{i}(0)\), \(s^{0}_{i}(0)=0\), \(r^{0}_{ij}(0)=0\) for \(i\neq j\). Then, for any fixed \(T\) (possibly dependent on \(m\) but not on \(d\)), the followings holds with probability at least \(1-\exp(-C^{\prime}m)\) over the i.i.d. initialization \((a_{i}(0),u_{i}(0))_{i\in[m]}\):_

\[\sup_{t\in[0,T]}|\mathscr{R}(a(t),u(t))-\mathscr{R}_{\mathrm{red} }(a^{0}(t),s^{0}(t),R^{0}(t))|\leq\frac{CM}{\sqrt{d}}\exp\left(MT(1+T)^{2}/ \varepsilon^{2}\right)\,, \tag{19}\] \[\max\left(\sup_{t\in[0,T]}\frac{1}{\sqrt{m}}\|a(t)-a^{0}(t)\|_{2 },\frac{1}{\sqrt{m}}\sup_{t\in[0,T]}\|s(t)-s^{0}(t)\|_{2}\right)\leq\frac{1}{ \sqrt{d}}\cdot C\exp\left(MT(1+T)^{2}/\varepsilon^{2}\right)\,,\] (20) \[\sup_{t\in[0,T]}\frac{1}{m}\|R(t)-R^{0}(t)\|_{\mathrm{F}}\leq\frac {1}{\sqrt{d}}\cdot C\exp\left(MT(1+T)^{2}/\varepsilon^{2}\right)\,. \tag{21}\]

_Here \(C,C^{\prime}\) are absolute constants and \(M\) only depends on the \(M_{i}\)'s in Assumptions A1-A3._

The proof of Corollary 1 is deferred to Appendix A.2. From now on, we will assume the initialization \(s^{0}_{i}(0)=0\), \(r^{0}_{ij}(0)=0\) for \(i\neq j\), but drop the superscript \(0\) for notational simplicity. We notice in passing that the right-hand sides of Eqs. (19) to (21) are independent of \(m\): this
approximation step holds uniformly over \(m\). (Note that the left hand sides are normalized by \(m\) as to yield the root mean square error per entry.)

In order to state the reduction \((ii)\) outlined above, we define the mean field risk as

\[\mathscr{R}_{\mbox{\tiny{\it mf}}}(a,s):=\mathscr{R}_{\mbox{\tiny{\it red}}}(a,s,R={s}{s}^{\top})=\frac{1}{2}\|\varphi\|_{L^{2}}^{2}-\frac{1}{m}\sum_{i=1}^{m}a _{i}V(s_{i})+\frac{1}{2m^{2}}\sum_{i,j=1}^{m}a_{i}a_{j}U(s_{i}s_{j})\,. \tag{22}\]

Further, we denote by \(\{a_{i}^{\mbox{\tiny{\it mf}}}(t),s_{i}^{\mbox{\tiny{\it mf}}}(t)\}_{i=1}^{m}\) the solution to the following ODEs:

\[\varepsilon\partial_{t}a_{i}= \,V(s_{i})-\frac{1}{m}\sum_{j=1}^{m}a_{j}U(s_{i}s_{j})\,, \tag{23}\] \[\partial_{t}s_{i}= \,a_{i}\left(1-s_{i}^{2}\right)\left(V^{\prime}(s_{i})-\frac{1}{ m}\sum_{j=1}^{m}a_{j}U^{\prime}(s_{i}s_{j})s_{j}\right)\,.\]

Note that (23) would be identical to (15)-(16) if we had \(r_{ij}=s_{i}s_{j}\). A priori, this is not the case. However, the two systems of equations are close to each other for large \(m\) as made precise by our next proposition, which formalizes reduction \((ii)\).

**Proposition 2** (Reduction to flow in \(\mathbb{R}^{2m}\)).: _Let \((a_{i}(t),s_{i}(t),r_{ij}(t))_{1\leq i<j\leq m}\) be the unique solution of the ODEs (15)-(18) with initialization \(s_{i}(0)=0\), \(r_{ij}(0)=0\) for all \(1\leq i\neq j\leq m\). Let \((a_{i}^{\mbox{\tiny{\it mf}}}(t),s_{i}^{\mbox{\tiny{\it mf}}}(t))_{i\leq m}\) be the unique solution of the ODEs (23) with initialization \(s_{i}^{\mbox{\tiny{\it mf}}}(0)=0\), \(a_{i}^{\mbox{\tiny{\it mf}}}(0)=a_{i}(0)\) for all \(i\leq m\)._

_If assumptions A1-A3 hold, then for any \(T<\infty\) there exists a constant_

\[C(T)=M\exp(MT(1+T)^{2}/\varepsilon^{2}) \tag{24}\]

_(with \(M\) depending on the constants \(\{M_{i}\}_{1\leq i\leq 3}\) appearing in Assumptions A1-A3 only) such that:_

\[\sup_{t\in[0,T]}\frac{1}{m}\sum_{i=1}^{m}\left\|(a_{i}(t),s_{i}(t))-(a_{i}^{ \mbox{\tiny{\it mf}}}(t),s_{i}^{\mbox{\tiny{\it mf}}}(t))\right\|_{2}^{2}\leq \frac{C(T)}{m}\,.\]

_Consequently,_

\[\sup_{t\in[0,T]}|\mathscr{R}_{\mbox{\tiny{\it red}}}\left(a(t),s(t),R(t) \right)-\mathscr{R}_{\mbox{\tiny{\it mf}}}\left(a^{\mbox{\tiny{\it mf}}}(t),s ^{\mbox{\tiny{\it mf}}}(t)\right)|\leq\frac{C(T)}{\sqrt{m}}\,.\]

The proof of this proposition is deferred to Appendix A.3. Now, combining the propositions and corollaries in this section, we deduce that with high probability over the i.i.d. initialization,

\[\sup_{t\in[0,T]}|\mathscr{R}(a(t),u(t))-\mathscr{R}_{\mbox{\tiny{\it mf}}} \left(a^{\mbox{\tiny{\it mf}}}(t),s^{\mbox{\tiny{\it mf}}}(t)\right)|\leq \left(\frac{1}{\sqrt{d}}+\frac{1}{\sqrt{m}}\right)CM\exp(MT(1+T)^{2}/ \varepsilon^{2}). \tag{25}\]

### Connection with mean field theory

Consider the empirical distributions of the neurons:

\[\widehat{\rho}_{t} :=\frac{1}{m}\sum_{i=1}^{m}\delta_{(a_{i}(t),s_{i}(t))}\,, \tag{26}\] \[\rho_{t} :=\frac{1}{m}\sum_{i=1}^{m}\delta_{(a_{i}^{\mbox{\tiny{\it mf}}} (t),s_{i}^{\mbox{\tiny{\it mf}}}(t))}\,, \tag{27}\]

with \((a_{i}(t),s_{i}(t))_{i\leq m}\), \((a_{i}^{\mbox{\tiny{\it mf}}}(t),s_{i}^{\mbox{\tiny{\it mf}}}(t))_{i\leq m}\) as in the statement of Proposition 2, i.e., solving (respectively) Eqs. (15)-(18) and Eq. (23) with initial conditions as given there.
Then, it is immediate to show that \(\rho_{t}\) solves (in weak sense) the following continuity partial differential equation (PDE) (we refer to Ambrosio et al. (2005); Santambrogio (2015) for the definition of weak solutions and basic properties, and Appendix A.4 for a short derivation.)

\[\partial_{t}\rho_{t}(a,s) =-\nabla\cdot\left(\rho_{t}\Psi\left(a,s;\rho_{t}\right)\right) \tag{28}\] \[:=-\left(\partial_{a}\left(\rho_{t}\Psi_{a}\left(a,s;\rho_{t} \right)\right)+\partial_{s}\left(\rho_{t}\Psi_{s}\left(a,s;\rho_{t}\right) \right)\right), \tag{29}\]

where \(\Psi=\left(\Psi_{a},\Psi_{s}\right)\) is given by

\[\Psi_{a}(a,s;\rho)= \,\varepsilon^{-1}\cdot\left(V(s)-\int_{\mathbb{R}^{2}}a_{1}U( ss_{1})\rho(\mathrm{d}a_{1},\mathrm{d}s_{1})\right), \tag{30}\] \[\Psi_{s}(a,s;\rho)= a(1-s^{2})\cdot\left(V^{\prime}(s)-\int_{\mathbb{R}^{2}}a_{1}s_{ 1}U^{\prime}(ss_{1})\rho(\mathrm{d}a_{1},\mathrm{d}s_{1})\right). \tag{31}\]

This equation can be extended to a flow in the whole space \((\mathscr{P}(\mathbb{R}^{2}),W_{2})\) (all probability measures on \(\mathbb{R}^{2}\) equipped with the second Wasserstein distance), and interpreted as gradient flow with respect to this metric in the following risk:

\[\mathscr{R}_{\text{mf},*}(\rho):=\frac{1}{2}\|\varphi\|_{L^{2}}^{2}-\int\!aV (s)\,\rho(\mathrm{d}a,\mathrm{d}s)+\frac{1}{2}\int\!a_{1}a_{2}U(s_{1}s_{2})\, \rho(\mathrm{d}a_{1},\mathrm{d}s_{1})\,\rho(\mathrm{d}a_{2},\mathrm{d}s_{2})\, \tag{32}\]

which is the obvious extension of \(\mathscr{R}_{\text{mf}}(a,s)\) of Eq. (22) to general probability distributions. Proposition 2 implies that for any \(T<\infty\), and under the above initial conditions,

\[\sup_{t\in[0,T]}W_{2}(\rho_{t},\widehat{\rho}_{t})\leq\sqrt{\frac{M\exp(MT(1+T )^{2}/\varepsilon^{2})}{m}}\,. \tag{33}\]

If we further denote by \(\rho_{t}^{d}\) the empirical distribution of \((a_{i}(t),s_{i}(t))\), \(i\leq m\), when \(s_{i}(0)=\langle u_{i}(0),u_{*}\rangle\), \(u_{i}(0)\sim\mathrm{Unif}(\mathbb{S}^{d-1})\), a further application of Corollary 1 yields

\[\sup_{t\in[0,T]}W_{2}(\rho_{t}^{d},\rho_{t})\leq\sqrt{\frac{M\exp(MT(1+T)^{2}/ \varepsilon^{2})}{m\wedge d}}\,. \tag{34}\]

Starting with Mei et al. (2018); Chizat and Bach (2018); Rotskoff and Vanden-Eijnden (2018), several authors used continuity PDEs of the form (28) to study the learning dynamics of two-layer neural networks. Following the physics tradition, this is referred to as the'mean-field theory' of two-layer neural networks. Appendix A.5 sketches an alternative approach to prove bounds of the form (25), (34) using the results of Mei et al. (2018, 2019). The present derivation has the advantages of yielding a sharper bound and of being self-contained.

### A general formulation

As mentioned above, the system of ODEs in Eq. (23) is a special case of the Wasserstein gradient flow of Eq. (28) whereby we set \(\rho_{0}=m^{-1}\sum_{i=1}^{m}\delta_{(a_{i}^{\text{mf}}(0),s_{i}^{\text{mf}}( 0))}\). In order to study the solutions of Eq. (28) (hence Eq. (23)) we adopt the following framework. Let \((\Omega,\rho)\) denote a probability space. Let \(a=a(\omega,t)\) and \(s=s(\omega,t)\) (\(\omega\in\Omega\), \(t\geqslant 0\)) be two measurable functions satisfying (dropping dependencies in \(t\) below)

\[\varepsilon\partial_{t}a(\omega)= \,V(s(\omega))-\int\mathrm{d}\rho(\nu)a(\nu)U(s(\omega)s(\nu))\,, \tag{35}\] \[\partial_{t}s(\omega)= \,a(\omega)\left(1-s(\omega)^{2}\right)\left(V^{\prime}(s(\omega) )-\int\mathrm{d}\rho(\nu)a(\nu)U^{\prime}(s(\omega)s(\nu))s(\nu)\right)\,.\]

If \(\omega=i\in\Omega=\{1,\dots,m\}\) endowed with the uniform measure, we obtain the equations (23). In general, the push-forward \(\rho_{t}\) of the measure \(\rho\) through the map \(\omega\in\Omega\mapsto(a(\omega,t),s(\omega,t))\in\mathbb{R}^{2}\) satisfies the mean-field equation (28). As a consequence, the dynamics (35) can be viewed as a gradient flow on the risk

\[\mathscr{R}_{\text{mf},*}(\rho)=\frac{1}{2}\|\varphi\|^{2}-\int a(\omega)V(s( \omega))\mathrm{d}\rho(\omega)+\frac{1}{2}\int a(\omega_{1})a(\omega_{2})U(s( \omega_{1})s(\omega_{2}))\mathrm{d}\rho(\omega_{1})\mathrm{d}\rho(\omega_{2} )\,. \tag{36}\]
## 5 Numerical solution

In Figure 2, we present the result of an Euler discretization of Eqs. (23) where \(\varphi\) is a degree-2 polynomial and \(\sigma\) is the ReLU activation: \(\sigma(s)=\max(s,0)\),

\[\begin{split}\varphi(s)&=\mathrm{He}_{0}(s)- \mathrm{He}_{1}(s)+\frac{2}{3}\mathrm{He}_{2}(s)\\ &=\left(1+\frac{2\sqrt{2}}{6}\right)-s+\frac{2\sqrt{2}}{6}s^{2} \,.\end{split} \tag{37}\]

These plots clearly display two of the features emphasized in the introduction: \((i)\) plateaus separated by periods of rapid improvement of the risk; \((ii)\) increasingly long timescales (notice the logarithmic time axis in the second and third row).

In order to examine the incremental learning structure, we rewrite the risk \(\mathscr{R}_{\mbox{\tiny mf}}\) of Eq. (22)

Figure 2: Simulation of the mean field neuron dynamics of Eqs. (23), with the target function of Eq. (37) and ReLU activations. We use learning rate ratios \(\varepsilon=10^{-3}\) (left) and \(\varepsilon=10^{-6}\) (right) and we use \(m=10\) neurons. First two rows: evolution of the risk \(\mathscr{R}_{\mbox{\tiny mf}}\) of Eq. (22), in linear and log-scales. Third row: evolution of the first three terms of the sum of (38).
by decomposing \(\varphi\) and \(\sigma\) in the basis of Hermite polynomials

\[\mathscr{R}_{\text{mf}}(a,s)=\frac{1}{2}\sum_{k\geqslant 0}\left(\varphi_{k}- \frac{\sigma_{k}}{m}\sum_{i=1}^{m}a_{i}s_{i}^{k}\right)^{2}\,. \tag{38}\]

We observe that, for small \(\varepsilon\), the Hermite coefficients of \(\varphi\) are learned sequentially, in the order of their degree. When \(\varepsilon\) is sufficiently small (right plots), this incremental learning happens in well separated phases. The plateaus and waterfalls in the plots of \(\mathscr{R}_{\text{mf}}\) correspond to the network learning increasingly higher degree polynomials.

In Figure 3 we plot the evolution of the values of the \(a_{i}\) and \(s_{i}\), for \(i\in\{1,\dots,m\}\). We observe that the order of magnitude of the \(a_{i}\)'s and the \(s_{i}\)'s increases when passing through the different phases of the incremental learning process.

Altogether, the results of Figures 2 and 3 are consistent with the standard learning scenario up to level \(L=2\) as per Definition 1. While we conjecture that incremental learning also occurs for higher-order polynomials, we found this hard to observe in numerical simulations.

First, as predicted in Definition 1, the times at which the components are learned are closer on a logarithmic scale as the degree increases. It is therefore increasingly difficult to observe time scales corresponding to higher degrees.

Second, we expect there to be a choice of the initialization \((a_{i,\text{init}},u_{i,\text{init}})_{i\in[m]}\), activation and target function, for which not all the components of \(\varphi\) are actually learnt. We observed empirically that this happens easily for small \(m\).

## 6 Timescales hierarchy in the gradient flow dynamics

We are interested in the behavior of the solution of the ODEs (35), initialized from \(s(\omega,0)=0\) for all \(\omega\) (as per Proposition 2). The standard learning scenario of Definition 1 concerns the behavior of solutions for \(\varepsilon\to 0\). This type of questions can be addressed within the theory of dynamical systems using _singular perturbation theory_(Holmes, 2013) ('singular' refers to the fact that \(\varepsilon\) multiplies one of the highest-order derivatives).

As a side remark, we note that the system (35) can be seen as a slow-fast dynamical system, where the \(a(\omega)\)'s are the fast variables and the \(s(\omega)\)'s are the slow variables (Berglund, 2001). Formally, the time derivative of the \(a(\omega)\)'s is multiplied by a factor \((1/\varepsilon)\). From a dynamical systems perspective, the present case is made complicated because of a bifurcation when the \(s(\omega)\)'s become non-zero.

The standard learning scenario provides a detailed description of this bifurcation. We will motivate this scenario using a classical, but non-rigorous, technique of singular perturbation

Figure 3: Same simulation as in Figure 2 (b). In these plots, we show the evolution of the \(a_{i}\) and the \(s_{i}\) for \(i\in\{1,\dots,m\}\) following a discretization of Eqs. (23).
theory, called the _matched asymptotic expansion_[Holmes, 2013, Chapter 2]. This technique decomposes the approximation of the solution in several time scales on which a regular approximation holds. These time scales are traditionally called _layers_ in the literature; however, we avoid this terminology due to the potential confusion with the layers of the neural network.

We will work mainly using the Hermite representation of the dynamical ODEs (35), which we write down for the reader's convenience:

\[\begin{split}\varepsilon\partial_{t}a(\omega)&= \,\sum_{k=0}^{\infty}\sigma_{k}s(\omega)^{k}\left(\varphi_{k}-\sigma_{k}\int a (\nu)s(\nu)^{k}\mathrm{d}\rho(\nu)\right)\,,\\ \partial_{t}s(\omega)&=\,a(\omega)\left(1-s(\omega) ^{2}\right)\sum_{k=1}^{\infty}k\sigma_{k}s(\omega)^{k-1}\left(\varphi_{k}- \sigma_{k}\int a(\nu)s(\nu)^{k}\mathrm{d}\rho(\nu)\right)\,.\end{split} \tag{39}\]

Sections 6.1-6.3 respectively describe the first three time scales of the matched asymptotic expansion of (39). This gives, for each time scale, an approximation of the \(a(\omega)\), \(s(\omega)\). In Appendix B.2, we detail how these sections induce an evolution of the risk alternating plateaus and rapid decreases, and support the standing learning scenario of Definition 1. Finally, in Section 6.4, we conjecture the behavior on longer time scales.

Notations.We denote \(\mathds{1}\) the constant function \(\mathds{1}:\omega\in\Omega\mapsto 1\in\mathbb{R}\). Denote \(\langle.,.\rangle_{L^{2}(\rho)}\) the dot product on \(L^{2}(\rho)\) and \(\|.\|_{L^{2}(\rho)}\) the associated norm. For \(x\in L^{2}(\rho)\), we denote \(x_{\perp}\) the orthogonal projection of \(x\) on the hyperplane \(\mathds{1}^{\perp}\) of \(L^{2}(\rho)\) of functions orthogonal to \(\mathds{1}\):

\[x_{\perp}(\omega)=x(\omega)-\int x(\nu)\mathrm{d}\rho(\nu)\,.\]

We denote \(a_{\mathrm{init}}(\omega)=a(\omega,0)\) and thus \(a_{\perp,\mathrm{init}}\) is the orthogonal projection of \(a_{\mathrm{init}}\) on \(\mathds{1}^{\perp}\).

### First time scale: constant component

We define a "fast" time variable \(t_{1}=t/\varepsilon\) and replace it in Eq. (39). We expand the solutions \(a(\omega)\) and \(s(\omega)\) in powers of \(\varepsilon\):

\[a(\omega) =a^{(0)}(\omega)+\varepsilon a^{(1)}(\omega)+\varepsilon^{2}a^{ (2)}(\omega)+\dots\,, \tag{40}\] \[s(\omega) =s^{(0)}(\omega)+\varepsilon s^{(1)}(\omega)+\varepsilon^{2}s^{ (2)}(\omega)+\dots\,, \tag{41}\]

where \(a^{(0)}(\omega),a^{(1)}(\omega),a^{(2)}(\omega),\dots,s^{(0)}(\omega),s^{(1)} (\omega),s^{(2)}(\omega),\dots\) are implicitly functions of \(t_{1}\). They are initialized at

\[a^{(0)}(\omega,t_{1}=0) =a_{\mathrm{init}}(\omega)\,, a^{(1)}(\omega,t_{1}=0) =0\,, a^{(2)}(\omega,t_{1}=0) =0\,, \dots \tag{42}\] \[s^{(0)}(\omega,t_{1}=0) =0\,, s^{(1)}(\omega,t_{1}=0) =0\,, s^{(2)}(\omega,t_{1}=0) =0\,, \dots \tag{43}\]

to be consistent with the initial condition \(a(\omega,t_{1}=0)=a(\omega,t=0)=a_{\mathrm{init}}(\omega)\) and \(s(\omega,t_{1}=0)=s(\omega,t=0)=0\).

We substitute the expansion in (39):

\[\partial_{t_{1}}a^{(0)}(\omega)+\varepsilon\partial_{t_{1}}a^{ (1)}(\omega)+\dots \tag{44}\] \[\quad=\sum_{k=0}^{\infty}\sigma_{k}\left(s^{(0)}(\omega)+ \varepsilon s^{(1)}(\omega)+\dots\right)^{k}\] (45) \[\qquad\times\left(\varphi_{k}-\sigma_{k}\int\left(a^{(0)}(\nu)+ \varepsilon a^{(1)}(\nu)+\dots\right)\left(s^{(0)}(\nu)+\varepsilon s^{(1)}( \nu)+\dots\right)^{k}\mathrm{d}\rho(\nu)\right)\,,\] (46) \[\partial_{t_{1}}s^{(0)}(\omega)+\varepsilon\partial_{t_{1}}s^{ (1)}(\omega)+\dots \tag{47}\]
\[=\varepsilon\left(a^{(0)}(\omega)+\varepsilon a^{(1)}(\omega)+\dots \right)\left(1-\left(s^{(0)}(\omega)+\varepsilon s^{(1)}(\omega)+\dots\right)^{2} \right)\sum_{k=1}^{\infty}k\sigma_{k}\left(s^{(0)}(\omega)+\varepsilon s^{(1)}( \omega)+\dots\right)^{k-1} \tag{48}\] \[\qquad\times\left(\varphi_{k}-\sigma_{k}\int\left(a^{(0)}(\nu)+ \varepsilon a^{(1)}(\nu)+\dots\right)\left(s^{(0)}(\nu)+\varepsilon s^{(1)}( \nu)+\dots\right)^{k}\mathrm{d}\rho(\nu)\right)\,. \tag{49}\]

The basic assumption of matched asymptotic expansions is that terms of the same order in \(\varepsilon\) can be identified (with some limitations that we develop below). For now, let us identify terms of order \(1=\varepsilon^{0}\):

\[\partial_{t_{1}}a^{(0)}(\omega) =\sum_{k=0}^{\infty}\sigma_{k}\left(s^{(0)}(\omega)\right)^{k} \left(\varphi_{k}-\sigma_{k}\int a^{(0)}(\nu)\left(s^{(0)}(\nu)\right)^{k} \mathrm{d}\rho(\nu)\right)\,, \tag{50}\] \[\partial_{t_{1}}s^{(0)}(\nu) =0\,. \tag{51}\]

From (51) and (43), we have \(s^{(0)}(\omega)=0\): time \(t_{1}=O(1)\Leftrightarrow t=O(\varepsilon)\) is too short for the \(s(\omega)\) to be of order \(1\).

Substituting \(s^{(0)}(\omega)=0\) in (50), we obtain

\[\partial_{t_{1}}a^{(0)}(\omega)=\sigma_{0}\left(\varphi_{0}-\sigma_{0}\int a^ {(0)}(\nu)\mathrm{d}\rho(\nu)\right)\,. \tag{52}\]

Recall that \(\langle.,.\rangle_{L^{2}(\rho)}\) is the dot product on \(L^{2}(\rho)\), \(\mathds{1}\) denotes the constant function \(\mathds{1}:\omega\in\Omega\mapsto 1\in\mathbb{R}\) and \(a_{\perp}\) is the orthogonal projection of \(a\) on \(\mathds{1}^{\perp}\). Equation (52) can be rewritten as

\[\partial_{t_{1}}\langle a^{(0)},\mathds{1}\rangle_{L^{2}(\rho)} =\sigma_{0}\left(\varphi_{0}-\sigma_{0}\langle a^{(0)},\mathds{1 }\rangle_{L^{2}(\rho)}\right)\,,\] \[\partial_{t_{1}}a^{(0)}_{\perp} =0\,,\]

which gives after integration (using (42)):

\[\langle a^{(0)},\mathds{1}\rangle_{L^{2}(\rho)} =e^{-\sigma_{0}^{2}t_{1}}(a_{\mathrm{init}},\mathds{1})_{L^{2}( \rho)}+\left(1-e^{-\sigma_{0}^{2}t_{1}}\right)\frac{\varphi_{0}}{\sigma_{0}}\,, \tag{53}\] \[a^{(0)}_{\perp} =a_{\perp,\mathrm{init}}\,.\]

At this point, we have determined \(a^{(0)}(\omega)\) and \(s^{(0)}(\omega)\), and thus \(a(\omega)=a^{(0)}(\omega)+O(\varepsilon)\) and \(s(\omega)=s^{(0)}(\omega)+O(\varepsilon)\) up to a \(O(\varepsilon)\) precision, which is sufficient to obtain a \(o(1)\)-approximation of the risk \(\mathscr{R}_{\mbox{\tiny mf,*}}\) (see Section B.2). However, note that we could obtain more precise estimates by identifying higher-order terms in (44)-(49). For instance, identifying the \(O(\varepsilon)\) terms in (47)-(49), we obtain \(\partial_{t_{1}}s^{(1)}(\omega)=a^{(0)}(\omega)\sigma_{1}\varphi_{1}\). This shows that the \(s(\omega)\) become non-zero, though only of order \(\varepsilon\) on the time scale \(t_{1}\asymp 1\); the inner-layer weights develop an infinitesimal correlation with the true direction \(u_{*}\) thanks to the linear component of \(\sigma\) and \(\varphi\).

The approximation constructed above should be considered as valid on the time scale \(t_{1}\asymp 1\Leftrightarrow t\asymp\varepsilon\). The approximation breaks down when we reach a new time scale, at which the \(s(\omega)\) are large enough for the \(a(\omega)\) to be affected (at leading order) by the linear part of the functions. We detail the new time scale and its resolution in the next section.

### Second time scale: linear component I

In this section, we seek a second, slower time scale, for which the behavior of the asymptotic expansion is different.
Identification of the scale.Consider \(t_{2}=\frac{t}{\varepsilon^{\gamma}}\), where \(\gamma<1\) is to be determined. We rewrite the system (39) using \(t_{2}\), and expand the solutions \(a(\omega)\) and \(s(\omega)\):

\[a(\omega) =a^{(0)}(\omega)+\varepsilon^{\delta}a^{(1)}(\omega)+\varepsilon^ {2\delta}a^{(2)}(\omega)+\ldots\,, \tag{54}\] \[s(\omega) =\varepsilon^{\delta}s^{(1)}(\omega)+\varepsilon^{2\delta}s^{(2) }(\omega)+\ldots\,. \tag{55}\]

(Since within the previous time scale we obtained \(s(\omega)=O(\varepsilon)\), it is natural to assume \(s^{(0)}(\omega)=0\).)

Let us pause to comment on our method.

Similarly to what has been done in the previous time scale, we will substitute the expansions (54)-(55) in the equations (39) in order to compute the different terms in the expansion. However, this step also allows us to compute the exponents \(\gamma\) and \(\delta\), that give respectively the new time scale and the size of the \(s(\omega)\)'s.

Note that we should have proceeded similarly for the first time scale, by introducing a first time variable \(t_{1}=\frac{t}{\varepsilon^{\gamma}}\), expanding \(a(\omega),s(\omega)\) in powers \(1,\varepsilon^{\delta^{\prime}},\varepsilon^{2\delta^{\prime}},\ldots\), and determining \(\gamma^{\prime}\) and \(\delta^{\prime}\) a posteriori. This would have led, indeed, to \(\gamma^{\prime}=1\) and \(\delta^{\prime}=1\). However, for simplicity, we preferred to fix these values that are natural a priori.

Finally, note that the expansions (40)-(41) and (54)-(55) are different, because they are valid on different time scales. In fact, the only coherence conditions that we require below is that the expansions match in a joint asymptotic where \(t_{1}=\frac{t}{\varepsilon}\to\infty\) and \(t_{2}=\frac{t}{\varepsilon^{\gamma}}\to 0\). We thus build different approximations for each one of the time scales, with some matching conditions; this justifies the name of _matched asymptotic expansion_.

We now return to our computations and substitute (54)-(55) in (39):

\[\varepsilon^{1-\gamma}\partial_{t_{2}}a^{(0)}(\omega)+\ldots =\sum_{k=0}^{\infty}\sigma_{k}\left(\varepsilon^{\delta}s^{(1)}( \omega)+\ldots\right)^{k}\] \[\qquad\qquad\qquad\times\left(\varphi_{k}-\sigma_{k}\int\left(a^ {(0)}(\nu)+\ldots\right)\left(\varepsilon^{\delta}s^{(1)}(\nu)+\ldots\right) ^{k}\mathrm{d}\rho(\nu)\right)\,,\] \[\varepsilon^{\delta}\partial_{t_{2}}s^{(1)}(\omega)+\ldots =\varepsilon^{\gamma}\left(a^{(0)}(\omega)+\ldots\right)\left(1- \left(\varepsilon^{\delta}s^{(1)}(\omega)+\ldots\right)^{2}\right)\sum_{k=1}^ {\infty}k\sigma_{k}\left(\varepsilon^{\delta}s^{(1)}(\omega)+\ldots\right)^{k-1}\] \[\qquad\qquad\qquad\qquad\times\left(\varphi_{k}-\sigma_{k}\int \left(a^{(0)}(\nu)+\ldots\right)\left(\varepsilon^{\delta}s^{(1)}(\nu)+\ldots \right)^{k}\mathrm{d}\rho(\nu)\right)\,,\]

and thus

\[\varepsilon^{1-\gamma}\partial_{t_{2}}a^{(0)}(\omega)+O( \varepsilon^{1-\gamma+\delta}) =\sigma_{0}\left(\varphi_{0}-\sigma_{0}\int a^{(0)}(\nu) \mathrm{d}\rho(\nu)\right) \tag{56}\] \[\qquad-\varepsilon^{\delta}\sigma_{0}^{2}\int a^{(1)}(\nu) \mathrm{d}\rho(\nu)+\varepsilon^{\delta}\sigma_{1}\varphi_{1}s^{(1)}(\omega) +O(\varepsilon^{2\delta})\,,\] (57) \[\varepsilon^{\delta}\partial_{t_{2}}s^{(1)}(\omega)+O(\varepsilon ^{2\delta}) =\varepsilon^{\gamma}\sigma_{1}\varphi_{1}a^{(0)}(\omega)+O( \varepsilon^{\gamma+\delta})\,. \tag{58}\]

For the first time scale, we chose \(\gamma=\delta=1\), so that the terms of order \(\varepsilon^{\delta}\) were negligible compared to \(\varepsilon^{1-\gamma}\partial_{t_{2}}a^{(0)}(\omega)\) in (56). This means that the linear components \(\sigma_{1},\varphi_{1}\) of the functions had no effect on the \(a(\omega)\) at leading order. We are now interested in a new time scale where \(\varepsilon^{1-\gamma}\partial_{t_{2}}a^{(0)}(\omega)\) and \(\varepsilon^{\delta}\sigma_{1}\varphi_{1}s^{(1)}(\omega)\) are of the same order, i.e., \(1-\gamma=\delta\); then the linear components play a role in the dynamics.

Further, for \(s^{(1)}(\omega)\) to be non-zero, we need both sides of (58) to be of the same order, thus \(\delta=\gamma\). Putting together, this gives \(\gamma=\delta=1/2\).
Derivation of the ODEs for this time scale.Let us summarize equations. For \(t_{2}=\frac{t}{\varepsilon^{\nicefrac{{1}}{{2}}}}\) and

\[a(\omega) =a^{(0)}(\omega)+\varepsilon^{\nicefrac{{1}}{{2}}}a^{(1)}(\omega)+ \ldots\,,\] \[s(\omega) =\varepsilon^{\nicefrac{{1}}{{2}}}s^{(1)}(\omega)+\ldots\,,\]

we have from (56)-(58):

\[\varepsilon^{\nicefrac{{1}}{{2}}}\partial_{t_{2}}a^{(0)}(\omega) =\sigma_{0}\left(\varphi_{0}-\sigma_{0}\int a^{(0)}(\nu)\mathrm{d} \rho(\nu)\right) \tag{59}\] \[\qquad-\varepsilon^{\nicefrac{{1}}{{2}}}\sigma_{0}^{2}\int a^{(1 )}(\nu)\mathrm{d}\rho(\nu)+\varepsilon^{\nicefrac{{1}}{{2}}}\sigma_{1}\varphi_ {1}s^{(1)}(\omega)+O(\varepsilon)\,,\] (60) \[\varepsilon^{\nicefrac{{1}}{{2}}}\partial_{t_{2}}s^{(1)}(\omega) =\varepsilon^{\nicefrac{{1}}{{2}}}\sigma_{1}\varphi_{1}a^{(0)}( \omega)+O(\varepsilon)\,. \tag{61}\]

First, we identify the terms of order \(1=\varepsilon^{0}\):

\[0=\sigma_{0}\left(\varphi_{0}-\sigma_{0}\int a^{(0)}(\nu)\mathrm{d}\rho(\nu) \right)\,. \tag{62}\]

This means that the trajectory remains in the affine hyperplane such that \(\varphi_{0}=\sigma_{0}\int a^{(0)}(\nu)\mathrm{d}\rho(\nu)\); intuitively, that the constant part of \(\varphi\) remains learned in this second time scale.

Second, we identify the terms of order \(\varepsilon^{\nicefrac{{1}}{{2}}}\) in (59)-(61):

\[\partial_{t_{2}}a^{(0)}(\omega) =-\sigma_{0}^{2}\int a^{(1)}(\nu)\mathrm{d}\rho(\nu)+\sigma_{1} \varphi_{1}s^{(1)}(\omega)\,, \tag{63}\] \[\partial_{t_{2}}s^{(1)}(\omega) =\sigma_{1}\varphi_{1}a^{(0)}(\omega)\,. \tag{64}\]

In (63), the first term of the right hand side depends on the unknown higher-order terms \(a^{(1)}(\nu)\); in fact, this is best interpreted as the Lagrange multiplier associated to the constraint (62). To eliminate this Lagrange multiplier, we use again the compact notations:

\[\partial_{t_{2}}a^{(0)} =-\sigma_{0}^{2}\langle a^{(1)},\mathds{1}\rangle_{L^{2}(\rho)} \mathds{1}+\sigma_{1}\varphi_{1}s^{(1)}\,, \tag{65}\] \[\partial_{t_{2}}s^{(1)} =\sigma_{1}\varphi_{1}a^{(0)}\,, \tag{66}\]

and thus

\[\partial_{t_{2}}a^{(0)}_{\perp} =\sigma_{1}\varphi_{1}s^{(1)}_{\perp}\,, \tag{67}\] \[\partial_{t_{2}}s^{(1)}_{\perp} =\sigma_{1}\varphi_{1}a^{(0)}_{\perp}\,. \tag{68}\]

Matching.The initialization of the ODEs (65)-(66) for the second time scale is determined by a classical procedure that matches with the previous time scale. In this paragraph, we denote \(\underline{a},\underline{s}\) the approximation obtained in the first time scale (Section 6.1), and \(\overline{a},\overline{s}\) the approximation in the second time scale, described above.

Consider an intermediate time scale \(\widetilde{t}=\frac{t}{\varepsilon^{\nicefrac{{1}}{{2}}}},\nicefrac{{1}}{{2 }}<\alpha<1\), and assume \(\widetilde{t}\asymp 1\) so that

\[t_{1} =\frac{t}{\varepsilon}=\frac{\widetilde{t}}{\varepsilon^{1- \alpha}}\to\infty\,, t_{2} =\frac{t}{\varepsilon^{\nicefrac{{1}}{{2}}}}=\varepsilon^{\alpha \nicefrac{{-1}}{{2}}}\widetilde{t}\to 0\,.\]

In this intermediate regime, we want the approximations provided on the first and the second time scales to match: \(\underline{a}(\widetilde{t})\) and \(\overline{a}(\widetilde{t})\) (resp. \(\underline{s}(\widetilde{t})\) and \(\overline{s}(\widetilde{t})\)) should match to leading order.

From the first time scale approximation,

\[\underline{a}=\underline{a}^{(0)}+O(\varepsilon) \tag{69}\]
\[= \langle\underline{a}^{(0)},\mathds{1}\rangle_{L^{2}(\rho)}\mathds{1}+ \underline{a}_{\perp}^{(0)}+O(\varepsilon) \tag{70}\] \[= \left[e^{-\sigma_{0}^{2}t_{1}}\langle\text{$\mathfrak{a}_{\text{ init}}$},\mathds{1}\rangle_{L^{2}(\rho)}+\left(1-e^{-\sigma_{0}^{2}t_{1}}\right) \frac{\varphi_{0}}{\sigma_{0}}\right]\mathds{1}+a_{\perp,\text{init}}+O(\varepsilon)\] (71) \[= \left[e^{-\sigma_{0}^{2}\widetilde{\gamma}/\varepsilon^{1-\alpha} }\langle a_{\text{init}},\mathds{1}\rangle_{L^{2}(\rho)}+\left(1-e^{-\sigma_{0 }^{2}\widetilde{\gamma}/\varepsilon^{1-\alpha}}\right)\frac{\varphi_{0}}{ \sigma_{0}}\right]\mathds{1}+a_{\perp,\text{init}}+O(\varepsilon)\] (72) \[= \frac{\varphi_{0}}{\sigma_{0}}\mathds{1}+a_{\perp,\text{init}}+o( 1)\,. \tag{73}\]

From the second time scale approximation,

\[\overline{a} =\overline{a}^{(0)}(t_{2})+O(\varepsilon^{\nicefrac{{1}}{{2}}} )=\overline{a}^{(0)}(\varepsilon^{\alpha-\nicefrac{{1}}{{2}}}\widetilde{t}) +O(\varepsilon^{\nicefrac{{1}}{{2}}}) \tag{74}\] \[=\overline{a}^{(0)}(0)+o(1)\,. \tag{75}\]

By matching, Equations (73) and (75) should be coherent. Thus the ODE for the second time scale should be initialized from \(\overline{a}^{(0)}(0)=\frac{\varphi_{0}}{\sigma_{0}}\mathds{1}+a_{\perp, \text{init}}\).

Similarly, the matching procedure gives that the ODE for the second time scale should be initialized from \(\overline{s}^{(1)}=0\).

Solution.As we are done with the matching procedure, we now consider the solution in the second time scale only, that we denote again by \(a\), \(s\) as in (65), (66). The matching procedure motivates us to consider the solution of (67)-(68) initialized at \(a_{\perp}^{(0)}(0)=a_{\perp,\text{init}}\), \(s_{\perp}^{(1)}=0\). This gives

\[a_{\perp}^{(0)} =\cosh\left(\varphi_{1}\sigma_{1}t_{2}\right)a_{\perp,\text{init }}\,, s_{\perp}^{(1)} =\sinh\left(\varphi_{1}\sigma_{1}t_{2}\right)a_{\perp,\text{init }}\,.\]

To conclude, we note that \(\langle a^{(0)},\mathds{1}\rangle_{L^{2}(\rho)}=\frac{\varphi_{0}}{\sigma_{0}}\) is constrained by (62). Further, from (64),

\[\partial_{t_{2}}\langle s^{(1)},\mathds{1}\rangle_{L^{2}(\rho)} =\sigma_{1}\varphi_{1}\langle a^{(0)},\mathds{1}\rangle_{L^{2}( \rho)}=\sigma_{1}\varphi_{1}\frac{\varphi_{0}}{\sigma_{0}},\]

thus \(\langle s^{(1)},\mathds{1}\rangle_{L^{2}(\rho)}=\sigma_{1}\varphi_{1}\frac{ \varphi_{0}}{\sigma_{0}}t_{2}\).

Putting together, these equations give:

\[a^{(0)} =\frac{\varphi_{0}}{\sigma_{0}}\mathds{1}+\cosh\left(\varphi_{1} \sigma_{1}t_{2}\right)a_{\perp,\text{init}}\,, s^{(1)} =\sigma_{1}\varphi_{1}\frac{\varphi_{0}}{\sigma_{0}}t_{2}\mathds{1}+\sinh \left(\varphi_{1}\sigma_{1}t_{2}\right)a_{\perp,\text{init}}\,. \tag{76}\]

We observe that \(a^{(0)}\) and \(s^{(1)}\) diverge as \(t_{2}\to\infty\). This implies that our approximation on the second time scale must break down at a certain point. Indeed, we analyzed this time scale under the assumption that both \(a^{(0)}\) and \(s^{(1)}\) are of order \(1\). However, since \(a^{(0)}\) and \(s^{(1)}\) diverge exponentially as \(t_{2}\to\infty\), as per Eq. (76), this assumption breaks down when \(t_{2}\asymp\log(1/\varepsilon)\).

More precisely, in (59) (resp. (61)), the \(O(\varepsilon)\) term includes a term of the form

\[-\varepsilon s^{(1)}(\omega)\sigma_{1}^{2}\int a^{(0)}(\nu)s^{(1)}(\nu)\text{d} \rho(\nu)\qquad\left(\text{resp. }-\varepsilon a^{(0)}(\omega)\sigma_{1}^{2}\int a^{(0)}(\nu)s^{(1)}(\nu)\text{d }\rho(\nu)\right)\,.\]

When \(a^{(0)}\) and \(s^{(1)}\) become of order \(\varepsilon^{-\nicefrac{{1}}{{4}}}\), this term becomes of order \(\varepsilon^{\nicefrac{{1}}{{4}}}\), which is then of the same order as the term \(\varepsilon^{\nicefrac{{1}}{{2}}}\sigma_{1}\varphi_{1}s^{(1)}(\omega)\) in (59) (resp. the term \(\varepsilon^{\nicefrac{{1}}{{2}}}\sigma_{1}\varphi_{1}a^{(0)}(\omega)\) in (61)). At this point, these terms can not be neglected anymore. From (76), we have

\[a^{(0)} \sim\frac{e^{\left|\varphi_{1}\sigma_{1}\right|t_{2}}}{2}a_{\perp, \text{init}}\,, s^{(1)} \sim\text{sign}(\varphi_{1}\sigma_{1})\frac{e^{\left|\varphi_{1}\sigma_{1} \right|t_{2}}}{2}a_{\perp,\text{init}}\,, t_{2}\to\infty\,.\]

Therefore, \(a^{(0)}\) and \(s^{(1)}\) become of order \(\varepsilon^{-\nicefrac{{1}}{{4}}}\) at the time \(t_{2}\sim\frac{1}{4\left|\sigma_{1}\varphi_{1}\right|}\log\frac{1}{\varepsilon}\), at which the approximation on the second time scale breaks down. We thus introduce a new time scale centered at this critical point.
### Third time scale: linear component II

We now introduce the time \(t_{3}=t_{2}-\frac{1}{4|\varphi_{1}\sigma_{1}|}\log\frac{1}{\varepsilon}\). As \(t_{3}\) is only a translation from \(t_{2}\), the ODEs in terms of \(t_{3}\) are the same as the ones in term of \(t_{2}\). However, in this time scale, \(a\) and \(\varepsilon^{\nicefrac{{1}}{{2}}}s\) have diverged. In coherence with the discussion above, we seek expansions of the form

\[a =\varepsilon^{-\nicefrac{{1}}{{4}}}a^{(-1)}+a^{(0)}+\varepsilon^{ \nicefrac{{1}}{{4}}}a^{(1)}+\ldots\,, \tag{77}\] \[s =\varepsilon^{\nicefrac{{1}}{{4}}}s^{(1)}+\varepsilon^{\nicefrac{{ 1}}{{2}}}s^{(2)}+\ldots\,. \tag{78}\]

Similarly to the second time scale, we substitute (77)-(78) in (39) and obtain

\[\varepsilon^{\nicefrac{{1}}{{4}}}\partial_{t_{3}}a^{(-1)}(\omega) =-\varepsilon^{-\nicefrac{{1}}{{4}}}\sigma_{0}^{2}\int a^{(-1)}( \nu)\mathrm{d}\rho(\nu)+\sigma_{0}\left(\varphi_{0}-\sigma_{0}\int a^{(0)}( \nu)\mathrm{d}\rho(\nu)\right)\] \[\quad-\varepsilon^{\nicefrac{{1}}{{4}}}\sigma_{0}^{2}\int a^{(1) }(\nu)\mathrm{d}\rho(\nu)+\varepsilon^{\nicefrac{{1}}{{4}}}\sigma_{1}\left( \varphi_{1}-\sigma_{1}\int a^{(-1)}(\nu)s^{(1)}(\nu)\mathrm{d}\rho(\nu) \right)s^{(1)}(\omega)+O(\varepsilon^{\nicefrac{{1}}{{2}}})\,,\] \[\varepsilon^{\nicefrac{{1}}{{4}}}\partial_{t_{3}}s^{(1)}(\omega) =\varepsilon^{\nicefrac{{1}}{{4}}}\sigma_{1}\left(\varphi_{1}- \sigma_{1}\int a^{(-1)}(\nu)s^{(1)}(\nu)\mathrm{d}\rho(\nu)\right)a^{(-1)}( \omega)+O(\varepsilon^{\nicefrac{{1}}{{2}}})\,.\]

First, we identify the terms of order \(\varepsilon^{-\nicefrac{{1}}{{4}}}\):

\[0=-\sigma_{0}^{2}\int a^{(-1)}(\nu)\mathrm{d}\rho(\nu)=-\sigma_{0}^{2}\left< a^{(-1)},\mathds{1}\right>_{L^{2}(\rho)}\,. \tag{79}\]

This means that \(a\) has no component diverging in \(\varepsilon\) in the direction of \(\mathds{1}\).

Second, we identify the terms of order \(1=\varepsilon^{0}\):

\[0=\sigma_{0}\left(\varphi_{0}-\sigma_{0}\int a^{(0)}(\nu)\mathrm{d}\rho(\nu) \right)=\sigma_{0}\left(\varphi_{0}-\sigma_{0}\left<a^{(0)},\mathds{1}\right> _{L^{2}(\rho)}\right)\,. \tag{80}\]

Put together with (79), this equation ensures that the constant component of \(\varphi\) remains learned on this third time scale.

Third, we identify the terms of order \(\varepsilon^{\nicefrac{{1}}{{4}}}\):

\[\partial_{t_{3}}a^{(-1)}(\omega) =-\sigma_{0}^{2}\int a^{(1)}(\nu)\mathrm{d}\rho(\nu)+\sigma_{1} \left(\varphi_{1}-\sigma_{1}\int a^{(-1)}(\nu)s^{(1)}(\nu)\mathrm{d}\rho(\nu) \right)s^{(1)}(\omega)\,, \tag{81}\] \[\partial_{t_{3}}s^{(1)}(\omega) =\sigma_{1}\left(\varphi_{1}-\sigma_{1}\int a^{(-1)}(\nu)s^{(1)}( \nu)\mathrm{d}\rho(\nu)\right)a^{(-1)}(\omega)\,.\]

Again, the term \(-\sigma_{0}^{2}\int a^{(1)}(\nu)\mathrm{d}\rho(\nu)\) is best interpreted as the Lagrange multiplier associated to the constraints (79), (80). Using the compact notations,

\[\int a^{(-1)}(\nu)s^{(1)}(\nu)\mathrm{d}\rho(\nu) =\left<a^{(-1)},s^{(1)}\right>_{L^{2}(\rho)}=\left<a^{(-1)}, \mathds{1}\right>_{L^{2}(\rho)}\left<\mathds{1},s^{(1)}\right>_{L^{2}(\rho)}+ \left<a^{(-1)}_{\perp},s^{(1)}_{\perp}\right>_{L^{2}(\rho)}\] \[=\left<a^{(-1)}_{\perp},s^{(1)}_{\perp}\right>_{L^{2}(\rho)}\,,\]

where in the last equality we use (79). Thus we can rewrite (81) as

\[\partial_{t_{3}}a^{(-1)} =-\sigma_{0}^{2}\langle a^{(1)},\mathds{1}\rangle_{L^{2}(\rho)} \mathds{1}+\sigma_{1}\left(\varphi_{1}-\sigma_{1}\left<a^{(-1)}_{\perp},s^{(1 )}_{\perp}\right>_{L^{2}(\rho)}\right)s^{(1)}\,, \tag{82}\] \[\partial_{t_{3}}s^{(1)} =\sigma_{1}\left(\varphi_{1}-\sigma_{1}\left<a^{(-1)}_{\perp},s^{(1 )}_{\perp}\right>_{L^{2}(\rho)}\right)a^{(-1)}\,,\]

and thus

\[\partial_{t_{3}}a^{(-1)}_{\perp} =\sigma_{1}\left(\varphi_{1}-\sigma_{1}\left<a^{(-1)}_{\perp},s^{( 1)}_{\perp}\right>_{L^{2}(\rho)}\right)s^{(1)}_{\perp}\,, \tag{83}\] \[\partial_{t_{3}}s^{(1)}_{\perp} =\sigma_{1}\left(\varphi_{1}-\sigma_{1}\left<a^{(-1)}_{\perp},s^{( 1)}_{\perp}\right>_{L^{2}(\rho)}\right)a^{(-1)}_{\perp}\,.\]
In Appendix B.1, we solve this system of ODEs and determine the initial condition by matching with the previous layer. The result is that

\[\begin{split}& a^{(-1)}=a_{\perp}^{(-1)}=\lambda a_{\perp,\text{ init}}\,,\\ & s^{(1)}=s_{\perp}^{(1)}=\text{sign}(\sigma_{1}\varphi_{1}) \lambda a_{\perp,\text{init}}\,,\end{split} \tag{84}\]

where \(\lambda=\lambda(t_{3})\) is the function

\[\lambda(t_{3})=\frac{|\varphi_{1}|^{\nicefrac{{1}}{{2}}}}{\left(|\sigma_{1}| \left\|a_{\perp,\text{init}}\right\|_{L^{2}(\rho)}^{2}+4|\varphi_{1}|e^{-2| \sigma_{1}\varphi_{1}|t_{3}}\right)^{\nicefrac{{1}}{{2}}}}\,. \tag{85}\]

This solution finishes to describe how the linear part of the function \(\varphi\) is learned.

### Conjectured behavior for larger time scales

The analysis of the previous sections naturally suggests the existence of a sequence of cutoffs. At each time scale, a new polynomial component of \(\varphi\) is learned within a window that is much shorter than the time elapsed before that phase started. Along this sequence, we expect \(s\) and \(a\) to grow to increasingly larger scales in \(\varepsilon\) (but \(s\) remains \(o(1)\) while \(a\) diverges).

More precisely, we assume that during the \(l\)-th phase, the network learns the degree-\(l\) component \(\varphi_{l}\), and various quantities satisfy the following scaling behavior:

\[a=O(\varepsilon^{-\omega_{l}}),\ \ \ s=O(\varepsilon^{\beta_{l}}),\ \ \ t=O( \varepsilon^{\mu_{l}})\,, \tag{86}\]

where \(\omega_{l}>0\) is an increasing sequence and \(\beta_{l},\mu_{l}>0\) are decreasing sequences. Further, while learning of this component takes place when \(t=O(\varepsilon^{\mu_{l}})\), the actual evolution of the risk (and of the neural network) take place on much shorter scales, namely:

\[\Delta t=O(\varepsilon^{\nu_{l}})\,, \tag{87}\]

where \(\nu_{l}\) is also decreasing, with \(\nu_{l}>\mu_{l}\). The goal of this section is to provide heuristic arguments to conjecture the values of \(\omega_{l}\), \(\beta_{l}\), \(\mu_{l}\) and \(\nu_{l}\). We will base this conjecture on a rigorous analysis of a simplified model.

The simplified model is motivated by the expectation (supported by the heuristics and simulations in the previous sections) that learning each component happens independently from the details of the evolution on previous time scales. In the simplified model, the activation function \(\sigma(x)\) is proportional to the \(l\)-th Hermite polynomial, namely \(\sigma(x)=\sigma_{l}\text{He}_{l}(x)\). This is the component of \(\sigma\) that we expect to be relevant on the \(l\)-th time scale. The gradient flow equations (39) then read:

\[\begin{split}\varepsilon\partial_{t}a(\omega)=&\ \sigma_{l}s(\omega)^{l}\left(\varphi_{l}-\sigma_{l}\int a(\nu)s(\nu)^{l}\text{ d}\rho(\nu)\right)\,,\\ \partial_{t}s(\omega)=&\ a(\omega)\left(1-s(\omega) ^{2}\right)l\sigma_{l}s(\omega)^{l-1}\left(\varphi_{l}-\sigma_{l}\int a(\nu) s(\nu)^{l}\text{d}\rho(\nu)\right)\,.\end{split} \tag{88}\]

with corresponding risk component

\[\mathscr{R}_{l}=\frac{1}{2}\left(\varphi_{l}-\sigma_{l}\int a(\nu)s(\nu)^{l} \text{d}\rho(\nu)\right)^{2}.\]

We capture the effect of learning dynamics on the previous time scales by the overall magnitude of the \(a(\omega)\)'s and \(s(\omega)\)'s at initialization. Namely, we choose the scale of initialization of the simplified model to be given by the end of the \((l-1)\)-th time scale, i.e., \(a(\omega)\asymp\varepsilon^{-\omega_{l-1}}\) and \(s(\omega)\asymp\varepsilon^{\beta_{l-1}}\). Further, in order for the \((l-1)\)-th component to be learned, namely

\[\int a(\nu)s(\nu)^{l-1}\text{d}\rho(\nu)\approx\frac{\varphi_{l-1}}{\sigma_{l- 1}}, \tag{89}\]
we require \(\omega_{l-1}=(l-1)\beta_{l-1}\) so that \(\int a(\nu)s(\nu)^{l-1}\mathrm{d}\rho(\nu)=\Theta(1)\). Analogously, we assume \(\omega_{l}=l\beta_{l}\).

Based on this consideration, we introduce the rescaled variables

\[\widetilde{a}(\omega)=\varepsilon^{\omega_{l}}a(\omega),\ \widetilde{s}(\omega)= \varepsilon^{-\beta_{l}}s(\omega),\ \text{where}\ \widetilde{a}(\omega,0)\asymp\varepsilon^{\omega_{l}-\omega_{l-1}}, \widetilde{s}(\omega,0)\asymp\varepsilon^{\beta_{l-1}-\beta_{l}}.\]

Rewriting Eq. (88) in terms of \(\widetilde{a}(\omega)\)'s and \(\widetilde{s}(\omega)\)'s, and using \(\omega_{l}=l\beta_{l}\), we get that

\[\begin{split}\varepsilon^{1-2l\beta_{l}}\partial_{l}\widetilde{a }(\omega)=&\,\sigma_{l}\widetilde{s}(\omega)^{l}\left(\varphi_{l }-\sigma_{l}\int\widetilde{a}(\nu)\widetilde{s}(\nu)^{l}\mathrm{d}\rho(\nu) \right)\\ \varepsilon^{2\beta_{l}}\partial_{l}\widetilde{s}(\omega)=& \,l\sigma_{l}\widetilde{a}(\omega)\widetilde{s}(\omega)^{l-1}\left(1- \varepsilon^{2\beta_{l}}\widetilde{s}(\omega)^{2}\right)\left(\varphi_{l}- \sigma_{l}\int\widetilde{a}(\nu)\widetilde{s}(\nu)^{l}\mathrm{d}\rho(\nu) \right).\end{split} \tag{90}\]

In order for the \(\widetilde{a}(\omega)\)'s and \(\widetilde{s}(\omega)\)'s to be learned simultaneously, we need \(1-2l\beta_{l}=2\beta_{l}\), which implies \(\beta_{l}=1/2(l+1)\). Making a further change of the time variable \(t=\varepsilon^{\nu_{l}}\tau\), where \(\nu_{l}=2\beta_{l}=1/(l+1)\), it follows that

\[\begin{split}\partial_{\tau}\widetilde{a}(\omega)=& \,\sigma_{l}\widetilde{s}(\omega)^{l}\left(\varphi_{l}-\sigma_{l} \int\widetilde{a}(\nu)\widetilde{s}(\nu)^{l}\mathrm{d}\rho(\nu)\right)\\ \partial_{\tau}\widetilde{s}(\omega)=&\,l\sigma_{l} \widetilde{a}(\omega)\widetilde{s}(\omega)^{l-1}\left(1-\varepsilon^{2\beta_{l }}\widetilde{s}(\omega)^{2}\right)\left(\varphi_{l}-\sigma_{l}\int\widetilde{a }(\nu)\widetilde{s}(\nu)^{l}\mathrm{d}\rho(\nu)\right).\end{split} \tag{91}\]

Moreover, rewriting the risk in terms of the rescaled variables \(\widetilde{a},\widetilde{s}\), \(\mathscr{R}_{l}(\tau)=\mathscr{R}_{l}(\widetilde{a}(\tau),\widetilde{s}(\tau))\) satisfies the ODE:

\[\partial_{\tau}\mathscr{R}_{l}=-2\sigma_{l}^{2}\mathscr{R}_{l}\cdot\int \widetilde{s}(\omega)^{2(l-1)}\left(l^{2}\widetilde{a}(\omega)^{2}\left(1- \varepsilon^{2\beta_{l}}\widetilde{s}(\omega)^{2}\right)+\widetilde{s}( \omega)^{2}\right)\mathrm{d}\rho(\omega). \tag{92}\]

Note that with our choice of \(\beta_{l}\) and \(\omega_{l}\), we have \(\omega_{l}-\omega_{l-1}=\beta_{l-1}-\beta_{l}=1/2l(l+1)\). This means that the \(\widetilde{a}(\omega)\)'s and \(\widetilde{s}(\omega)\)'s are initialized at the same scale, namely

\[\widetilde{a}(\omega,0),\widetilde{s}(\omega,0)=\Theta(\varepsilon^{1/2l(l+1) })\,. \tag{93}\]

The theorem below describes quantitatively the dynamics of the simplified model for small \(\varepsilon\), and determines the value of \(\mu_{l}\) (recall that \(\nu_{l}=1/(l+1)\)):

**Theorem 1** (Evolution of the simplified gradient flow).: _Assume \(l\geq 2\) and let \((\widetilde{a}(\omega,\tau),\widetilde{s}(\omega,\tau))_{\tau\geq 0}\) be the unique solution of the ODE system (91), initialized as per Eq. (93) (note in particular that \(\sigma_{l}\varphi_{l}\widetilde{a}(\omega,0)\widetilde{s}(\omega,0)^{l}\asymp \varepsilon^{1/2l}\)). Then the followings hold:_

* _Let us denote_ \[A=\left\{\omega:\sigma_{l}\varphi_{l}\liminf_{\varepsilon\to 0}\varepsilon^{-1/2l} \widetilde{a}(\omega,0)\widetilde{s}(\omega,0)^{l}>0\right\}\] (94) _and assume_ \(\rho(A)>0\)_. For_ \(\Delta\in(0,\varphi_{l}^{2}/2)\)_, define_ \[\tau(\Delta)=\inf\{\tau\geq 0:\mathscr{R}_{l}(\widetilde{a}(\tau), \widetilde{s}(\tau))\leq\Delta\}.\] (95) _Then, for any fixed_ \(\Delta\) _we have_ \(\tau(\Delta)=\Theta(\varepsilon^{-(l-1)/2l(l+1)})\) _as_ \(\varepsilon\to 0\)_. Further, if_ \(\rho\) _is a discrete probability measure, then there exists_ \(\tau_{*}(\varepsilon)=\Theta(\varepsilon^{-(l-1)/2l(l+1)})\) _and, for any_ \(\Delta>0\) _a constant_ \(c_{*}(\Delta)>0\) _independent of_ \(\varepsilon\) _such that_ \[\tau\leq\tau_{*}(\varepsilon)-c_{*}(\Delta)\Rightarrow \liminf_{\varepsilon\to 0}\mathscr{R}_{l}(\widetilde{a}(\tau), \widetilde{s}(\tau))\geq\frac{1}{2}\varphi_{l}^{2}-\Delta\,,\] (96) \[\tau\geq\tau_{*}(\varepsilon)+c_{*}(\Delta)\Rightarrow \limsup_{\varepsilon\to 0}\mathscr{R}_{l}(\widetilde{a}(\tau), \widetilde{s}(\tau))\leq\Delta\,,\] (97) _namely the_ \(l\)_-th component is learnt in an_ \(O(1)\) _time window around_ \(\tau_{*}(\varepsilon)=\Theta(\varepsilon^{-(l-1)/2l(l+1)})\)_._
_._
2. _Similarly, we denote_ \[B=\left\{\omega:\sigma_{l}\varphi_{l}\limsup_{\varepsilon\to 0}\varepsilon^{-1/2l} \widetilde{a}(\omega,0)\widetilde{s}(\omega,0)^{l}<0,\text{ and }\liminf_{\varepsilon\to 0}( \widetilde{s}(\omega,0)^{2}/\widetilde{a}(\omega,0)^{2})>l\right\}.\] _If_ \(\rho(B)>0\)_, then the same claims as in_ \((a)\) _hold._
3. _If neither of the conditions at points_ \((a)\)_,_ \((b)\) _holds, and_ \[\sigma_{l}\varphi_{l}\limsup_{\varepsilon\to 0}\varepsilon^{-1/2l} \widetilde{a}(\omega,0)\widetilde{s}(\omega,0)^{l}<0,\quad\limsup_{ \varepsilon\to 0}(\widetilde{s}(\omega,0)^{2}/\widetilde{a}(\omega,0)^{2})<l\] (99) _for almost every_ \(\omega\in\Omega\)_. Then, for such_ \(\omega\in\Omega\) _and each_ \(\Delta>0\)_, there exists a constant_ \(C_{*}(\omega,\Delta)>0\) _such that_ \[\tau\geq C_{*}(\omega,\Delta)\varepsilon^{-(l-1)/2l(l+1)}\ \Rightarrow\ | \widetilde{s}(\omega,\tau)|\leq\Delta\varepsilon^{1/2l(l+1)}\,,\] (100) _meaning that_ \(\widetilde{s}(\omega,\tau)\) _converges to_ \(0\) _eventually._

_We further note that \(\tau=\Theta(\varepsilon^{-(l-1)/2l(l+1)})\Longleftrightarrow t=\Theta( \varepsilon^{\mu_{l}})\) with \(\mu_{l}=1/2l\), and \(\tau=O(1)\Longleftrightarrow t=O(\varepsilon^{\mu_{l}})\) with \(\nu_{l}=1/(l+1)\)._

The proof of Theorem 1 is deferred to Appendix B.3.

**Remark 6.1**.: Under the conditions of cases \((a)\) and \((b)\), we see that the degree-\(l\) component of the target function is learnt within an \(O(\varepsilon^{1/(l+1)})\) time window around \(t_{*}(l,\varepsilon)\asymp\varepsilon^{1/2l}\), which is consistent with the timescales conjectured in Definition 1.

**Remark 6.2**.: Case \((c)\) corresponds to \(s(\omega)/s(\omega,0)\) becoming close to \(0\) in time \(t=O(\varepsilon^{\mu_{l}})\), and staying at \(0\). In other words, the neurons become orthogonal to the target direction and play no role in learning higher-degree components any longer.

Informally, case \((c)\) couples the learning of different polynomial components. It can happen that the learning phase \(l-1\) induces an effective initialization \((\widetilde{a}(\omega,0),\ \widetilde{s}(\omega,0))\) within the domain of case \((c)\).

We expect this not to be the case for suitable choices of initialization (or equivalently \(\mathrm{P}_{A}\)), \(\varphi\), and \(\sigma\). Establishing this would amount to establishing that the standard learning scenario holds.

## 7 Stochastic gradient descent and finite sample size

So far we focused on analyzing the projected gradient flow (GF) dynamics with respect to the population risk, as defined in Eqs. (4)-(5). In this section, we extract the implications of our analysis of GF on online projected stochastic gradient descent, which is a projected version of the SGD dynamics (151).

For simplicity of notation, we denote by \(z=(y,x)\in\mathbb{R}\times\mathbb{R}^{d}\) a datapoint and by \(\theta_{i}=(a_{i},u_{i})\in\mathbb{R}\times\mathbb{S}^{d-1}\) the parameters of neuron \(i\). For \(z=(y,x)\) and \(\rho^{(m)}=(1/m)\sum_{i=1}^{m}\delta_{\theta_{i}}=(1/m)\sum_{i=1}^{m}\delta_{( a_{i},u_{i})}\), we define

\[\widehat{F}_{i}(\rho^{(m)};z) =\left(y-\frac{1}{m}\sum_{j=1}^{m}a_{j}\sigma(\langle u_{j},x \rangle)\right)\sigma(\langle u_{i},x\rangle),\] \[\widehat{G}_{i}(\rho^{(m)};z) =a_{i}\left(y-\frac{1}{m}\sum_{j=1}^{m}a_{j}\sigma(\langle u_{j},x\rangle)\right)\sigma^{\prime}(\langle u_{i},x\rangle)x.\]

The projected SGD dynamics is specified as follows:

\[\overline{a}_{i}(k+1) =\overline{a}_{i}(k)+\varepsilon^{-1}\eta\widehat{F}_{i}( \overline{\rho}^{(m)}(k);z_{k+1}) \tag{101}\] \[\overline{u}_{i}(k+1) =\operatorname{Proj}_{\mathbb{S}^{d-1}}\left(\overline{u}_{i}(k) +\eta\widehat{G}_{i}(\overline{\rho}^{(m)}(k);z_{k+1})\right),\]
where for \(u\in\mathbb{R}^{d}\) and compact \(S\subset\mathbb{R}^{d}\), \(\text{Proj}_{S}(u):=\text{argmin}_{s\in S}\left\|s-u\right\|_{2}\), and \(\overline{\rho}^{(m)}:=(1/m)\sum_{i=1}^{m}\delta_{\overline{\theta}_{i}}\). Note that the \((\overline{a}_{i},\overline{u}_{i})\)'s here are different from the \((\overline{a},\overline{s})\)'s in Section 6.

We prove that, for small \(\eta\), the projected SGD of Eq. (101) is close to the gradient flow of Eqs. (4)-(5). Throughout this section, we make the following assumptions similar to those assumed in Section 4:

**A1.**: \(\rho_{0}\) is supported on \([-M_{1},M_{1}]\times\mathbb{S}^{d-1}\). Hence, \(\left|a_{i}(0)\right|\leq M_{1}\) for all \(i\in[m]\).
**A2.**: The activation function is bounded: \(\left\|\sigma\right\|_{\infty}\leq M_{2}\). Additionally, define for \(u,u^{\prime}\in\mathbb{R}^{d}\):

\[V(\left\langle u_{*},u\right\rangle;\left\|u_{*}\right\|_{2}, \left\|u\right\|_{2}) =\mathbb{E}\left[\varphi(\left\langle u_{*},x\right\rangle) \sigma(\left\langle u,x\right\rangle)\right], \tag{102}\] \[U(\left\langle u,u^{\prime}\right\rangle;\left\|u\right\|_{2}, \left\|u^{\prime}\right\|_{2}) =\mathbb{E}\left[\sigma(\left\langle u,x\right\rangle)\sigma( \left\langle u^{\prime},x\right\rangle)\right]. \tag{103}\]

We then require the functions \(V\) and \(U\) to be bounded and differentiable, with uniformly bounded and Lipschitz continuous gradients for all \(\left\|u\right\|_{2},\left\|u^{\prime}\right\|_{2}\leq 2\):

\[\left\|\nabla_{u}V\right\|_{2}\leq M_{2},\;\left\|\nabla_{u}V- \nabla_{u^{\prime}}V\right\|_{2}\leq M_{2}\left\|u-u^{\prime}\right\|_{2}, \tag{104}\] \[\left\|\nabla_{(u,u^{\prime})}U\right\|_{2}\leq M_{2},\;\left\| \nabla_{(u,u^{\prime})}U-\nabla_{(u_{1},u_{1}^{\prime})}U\right\|_{2}\leq M_{ 2}\left(\left\|u-u_{1}\right\|_{2}+\left\|u^{\prime}-u_{1}^{\prime}\right\|_{2 }\right). \tag{105}\]

Similar to Remark 4.1, we can show that a sufficient condition for Eq.s (104) and (105) is

\[\sup\left\{\left\|\sigma^{\prime}\right\|_{L^{2}},\left\|\sigma^{ \prime\prime}\right\|_{L^{2}}\right\}\leq M_{2}^{\prime},\quad\sup\left\{ \left\|\varphi\right\|_{L^{2}},\left\|\varphi^{\prime}\right\|_{L^{2}},\left\| \varphi^{\prime}\right\|_{L^{2}}\right\}\leq M_{2}^{\prime},\]

where the constant \(M_{2}^{\prime}\) depends uniquely on \(M_{2}\).
**A3.**: Assume \((x,y)\sim\mathds{P}\), then we require that \(y\in[-M_{3},M_{3}]\) almost surely. Moreover, we assume that for all \(\left\|u\right\|_{2}\leq 2\), both \(\sigma(\left\langle u,x\right\rangle)\) and \(\sigma^{\prime}(\left\langle u,x\right\rangle)(x-\left\langle u,x\right\rangle u)\) are \(M_{3}\)-sub-Gaussian.

The following theorem upper bounds the distance between gradient flow and projected stochastic gradient descent dynamics.

**Theorem 2** (Difference between GF and Projected SGD).: _Let \(\theta_{i}(t)=(a_{i}(t),u_{i}(t))\) be the solution of the GF ordinary differential equations (4)-(5). There exists a constant \(M\) that only depends on the \(M_{i}\)'s from Assumptions A1-A3, such that for any \(T,z\geq 0\) and_

\[\eta\leq\frac{1}{(d+\log m+z^{2})M\exp((1+1/\varepsilon)MT(1+T/ \varepsilon)^{2})},\]

_the following holds with probability at least \(1-\exp(-z^{2})\):_

\[\sup_{k\in[0,T/\eta]\cap\mathbb{N}}\max_{i\in[m]}|\overline{a}_{i} (k)| \leq M(1+T/\varepsilon), \tag{106}\] \[\sup_{k\in[0,T/\eta]\cap\mathbb{N}}\max_{i\in[m]}\left\|\theta_{i }(k\eta)-\overline{\theta}_{i}(k)\right\|_{2} \leq\left(\sqrt{d+\log m}+z\right)M\exp((1+1/\varepsilon)MT(1+T/ \varepsilon)^{2})\sqrt{\eta},\] (107) \[\sup_{k\in[0,T/\eta]\cap\mathbb{N}}|\mathscr{R}(\overline{a}(k), \overline{u}(k))-\mathscr{R}(a(k\eta),u(k\eta))| \leq\left(\sqrt{d+\log m}+z\right)M\exp((1+1/\varepsilon)MT(1+T/ \varepsilon)^{2})\sqrt{\eta}. \tag{108}\]

The proof is presented in Appendix C and follows the same scheme as in that of Theorem 1 part (B) in [11]. The main difference with respect to that theorem is here we are interested in projected SGD (and GF) instead of plain SGD (and GF), hence an additional step of approximation is required, and the \(a_{i}\)'s and \(u_{i}\)'s need to be treated separately. We next draw implications of the last result on learning by online SGD within the standard learning scenario.
**Theorem 3**.: _Fix any \(\delta>0\). Assume \(\varphi,\sigma\) and the initialization \(\mathrm{P}_{A}\) be such that the standard learning scenario of Definition 1 holds up to level \(L\) for some \(L\geq 2\), and that_

\[\sum_{k\geq L+1}\varphi_{k}^{2}\leq\frac{\delta}{2}. \tag{109}\]

_Then, there exist constants \(\varepsilon_{*}=\varepsilon_{*}(\delta)\), \(T_{0}=T_{0}(\delta)\), \(T=T(\varepsilon,\delta)=T_{0}(\delta)\varepsilon^{1/(2L)}\) and \(M=M(\varepsilon,\delta)\) that depend on \(\varepsilon,\delta\) (together with \(\varphi,\sigma\) and \(\mathrm{P}_{A}\)) such that the following happens. Assume \(\varepsilon\leq\varepsilon_{*}(\delta)\) and \(m,d,z\) are such that \(d\geq M\), \(m\geq\max(M,z)\), and the step size \(\eta\) and number of samples (equivalently, number of steps) \(n\) satisfy_

\[\eta =\frac{1}{M(d+\log m+z)}\,, \tag{110}\] \[n =MT(d+\log m+z)\,. \tag{111}\]

_Then, with probability at least \(1-e^{-z}\), the projected gradient descent algorithm of Eq. (101) achieves population risk smaller than \(\delta\):_

\[\mathds{P}\Big{(}\mathscr{R}(\overline{a}(n),\overline{u}(n))\leq\delta\Big{)} \geq 1-e^{-z}\,. \tag{112}\]

The proof of Theorem 3 is deferred to Appendix C.4.

**Remark 7.1**.: Within the lazy or neural tangent regime, learning the projection of the target function \(\varphi(\langle u_{*},x\rangle)\) onto polynomials of degree \(\ell\) requires \(n\gg d^{\ell}\) samples, and \(m\gg d^{\ell-1}\) neurons (Ghorbani et al., 2021, Mei et al., 2022, Montanari and Zhong, 2022).

In contrast, Theorem 3 shows that, within the standard learning scenario, \(O(d)\) samples and \(O(1)\) neurons are sufficient. Further as per Theorem 2, the learning dynamics is accurately described by the GF analyzed in the previous sections.

## Acknowledgments

This work was supported by the NSF through award DMS-2031883, the Simons Foundation through Award 814639 for the Collaboration on the Theoretical Foundations of Deep Learning, the NSF grant CCF-2006489 and the ONR grant N00014-18-1-2729, and a grant from Eric and Wendy Schmidt at the Institute for Advanced Studies. Part of this work was carried out while Andrea Montanari was on partial leave from Stanford and a Chief Scientist at Ndata Inc dba Project N. The present research is unrelated to AM's activity while on leave.