Learning time-scales in two-layers neural networks

Introduction

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

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

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., 2020a; Barak et al., 2022).

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.

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#2\#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#3\#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.

Setting and standard learning scenario

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., 2018b; Chizat and Bach, 2018; Ghorbani et al., 2020b; 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\leqslant i\leqslant m}$ on the population risk

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 $\mathscrsfs{R}(a,u)$ is defined by the following ordinary differential equations (ODEs):

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

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,\rm{init}},u_{i,\rm{init}})$:

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

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 ${\rm P}_{A}$ is square integrable)

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

We say that the standard learning scenario holds up to level $L$ for a certain target function $\varphi$, activation $\sigma$, and distribution ${\rm P}_{A}$, if the followings hold:

There exist constants $c_{2},\dots,c_{L+1}>0$ such that the following asymptotic holds as $\varepsilon\to 0$, $t\to 0$:

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,{\rm 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 the first Hermite $k+1$ coefficients of $\varphi$ vanish, $\varphi_{0}=\dots=\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.

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

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., 2018b; 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. (2018b) 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. (2018a) 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

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

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.

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. (2018b); Rotskoff and Vanden-Eijnden (2018); Chizat and Bach (2018); Mei et al. (2019) (with the required modifications for the constraints $\|u_{i}\|=1$).

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, ${\rm P}_{A}$ is supported on $[-M_{1},M_{1}]$.

The activation function is bounded: $\left\|{\sigma}\right\|_{\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$. A sufficient condition for this is

Responses are bounded, i.e., $\|\varphi\|_{\infty}\leq M_{3}$.

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\|{\sigma^{\prime}}\right\|_{L^{2}}\leq M_{2}$ and $\left\|{\varphi^{\prime}}\right\|_{L^{2}}\leq M_{2}$, then we have

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.

Define $s_{i}=\langle u_{i},u_{*}\rangle$, $r_{ij}=\langle u_{i},u_{j}\rangle$ for $i,j=1,\dots,m$. Then, letting $R=(r_{ij})_{i,j\leq m}$, we have

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)

This discussion immediately yields the following consequence.

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]}$:

Here $C,C^{\prime}$ are absolute constants and $M$ only depends on the $M_{i}$’s in Assumptions A1-A3.

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

Further, we denote by $\{a^{\mbox{\tiny\rm mf}}_{i}(t),s^{\mbox{\tiny\rm mf}}_{i}(t)\}_{i=1}^{m}$ the solution to the following ODEs:

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

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^{\mbox{\tiny\rm mf}}_{i}(t),s^{\mbox{\tiny\rm mf}}_{i}(t))_{i\leq m}$ be the unique solution of the ODEs (23) with initialization $s^{\mbox{\tiny\rm mf}}_{i}(0)=0$, $a^{\mbox{\tiny\rm mf}}_{i}(0)=a_{i}(0)$ for all $i\leq m$.

If assumptions A1-A3 hold, then for any $T<\infty$ there exists a constant

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

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,

Consider the empirical distributions of the neurons:

with $(a_{i}(t),s_{i}(t))_{i\leq m}$, $(a^{\mbox{\tiny\rm mf}}_{i}(t),s^{\mbox{\tiny\rm mf}}_{i}(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.)

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

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

Starting with Mei et al. (2018b); 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. (2018b, 2019). The present derivation has the advantages of yielding a sharper bound and of being self-contained.

2 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}^{\mbox{\tiny\rm mf}}(0),s_{i}^{\mbox{\tiny\rm 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)

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)$,

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 $\mathscrsfs{R}_{\mbox{\tiny\rm mf}}$ of Eq. (22) by decomposing $\varphi$ and $\sigma$ in the basis of Hermite polynomials

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 $\mathscrsfs{R}_{\mbox{\tiny\rm 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,\rm{init}},u_{i,\rm{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$.

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

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.

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

1 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$:

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

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

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}$:

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

which gives after integration (using (42)):

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 $\mathscrsfs{R}_{\mbox{\tiny\rm 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.

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

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)$:

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

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^{\prime}}}$, expanding $a(\omega),s(\omega)$ in powers $1,\varepsilon^{\delta^{\prime}},\varepsilon^{2\delta^{\prime}},\dots$, 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):

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

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

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

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:

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^{\alpha}}$, $\nicefrac{{1}}{{2}}<\alpha<1$, and assume $\widetilde{t}\asymp 1$ so that

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 second time scale approximation,

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,\rm{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,\rm{init}}$, $s_{\perp}^{(1)}=0$. This gives

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),

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

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

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

Therefore, $a^{(0)}$ and $s^{(1)}$ become of order $\varepsilon^{-\nicefrac{{1}}{{4}}}$ at the time $t_{2}\sim\frac{1}{4|\sigma_{1}\varphi_{1}|}\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.

3 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

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

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

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}$:

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}}}$:

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

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

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

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

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

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:

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.

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

Based on this consideration, we introduce the rescaled variables

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

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

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

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

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)$):

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:

and assume $\rho(A)>0$. For $\Delta\in(0,\varphi_{l}^{2}/2)$, define

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

namely the $l$-th component is learnt in an $O(1)$ time window around $\tau_{*}(\varepsilon)=\Theta(\varepsilon^{-(l-1)/2l(l+1)})$.

If $\rho(B)>0$, then the same claims as in $(a)$ hold.

If neither of the conditions at points $(a)$, $(b)$ holds, and

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

meaning that $\widetilde{s}(\omega,\tau)$ converges to 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^{\nu_{l}})$ with $\nu_{l}=1/(l+1)$.

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

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.

Case $(c)$ corresponds to $s(\omega)/s(\omega,0)$ becoming close to in time $t=O(\varepsilon^{\mu_{l}})$, and staying at . 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 ${\rm P}_{A}$), $\varphi$, and $\sigma$. Establishing this would amount to establishing that the standard learning scenario holds.

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

The projected SGD dynamics is specified as follows:

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:

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

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

where the constant $M_{2}^{\prime}$ depends uniquely on $M_{2}$.

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(\langle u,x\rangle)$ and $\sigma^{\prime}(\langle u,x\rangle)(x-\langle u,x\rangle u)$ are $M_{3}$-sub-Gaussian.

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

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

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

The proof is presented in Appendix C and follows the same scheme as in that of Theorem 1 part (B) in (Mei et al., 2019). 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.

Fix any $\delta>0$. Assume $\varphi,\sigma$ and the initialization ${\rm P}_{A}$ be such that the standard learning scenario of Definition 1 holds up to level $L$ for some $L\geq 2$, and that

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 ${\rm 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

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

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

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.

References

Appendix A Appendix to Section 4

This proves (14). Equation (15) follows directly:

To obtain equations (16)-(18), we now take gradients in (113):

This gives (16). Finally, we perform a similar computation to compute $\partial_{t}r_{ij}=\langle\partial_{t}u_{i},u_{j}\rangle+\langle u_{i},\partial_{t}u_{j}\rangle$. We compute only the first term, as the second term can be obtained by inverting $i$ and $j$:

Adding the symmetric term $\langle u_{i},\partial_{t}u_{j}\rangle$, we obtain (17)-(18).

A.2 Proof of Corollary 1

First, note that in the proof of Lemma 1, we obtain the following a priori estimate on the magnitude of the $a_{i}^{0}$’s:

where $M$ only depends on the $M_{i}$’s in Assumptions A1-A3. Using a similar argument as that in the proof of Proposition 2, we obtain that for any $t\in[0,T]$ and $i\in[m]$,

then we know that $G^{\prime}(t)\leq(M(1+t)^{2}/\varepsilon^{2})G(t)$. Applying Grönwall’s inequality yields

with probability at least $1-\exp(C^{\prime}m)$, where $C$ and $C^{\prime}$ are both absolute constants. Therefore,

Next we upper bound the risk difference, by direct calculation,

with probability at least $1-\exp(-C^{\prime}m)$, where the constant $M$ only depends on the $M_{i}$’s from Assumptions A1-A3. The conclusion now follows from taking the supremum over all $t\in[0,T]$. This completes the proof of Corollary 1.

A.3 Proof of Proposition 2

We consider $r_{ij}^{\perp}=r_{ij}-s_{i}s_{j}=\langle u_{i},u_{j}\rangle-\langle u_{i},u_{*}\rangle\langle u_{*},u_{j}\rangle$, the dot product between $u_{i}$ and $u_{j}$ that is out of the relevant subspace spanned by $u_{*}$. We show that these variables satisfy the ODEs

By definition of $r_{ij}^{\perp}$, we readily see that

If Assumptions A1-A3 hold, then we have for any fixed $T>0$:

To begin with, using Eq. (119), we obtain that

Using the ODEs for the $a_{i}$’s, we obtain that

where $(i)$ follows from our assumptions and the fact that $\mathscrsfs{R}(a(t),u(t))\leq\mathscrsfs{R}(a(0),u(0))$, since $\partial_{t}\mathscrsfs{R}(a,u)\leq 0$ by gradient flow equations. Moreover, the constant $M$ only depends on the $M_{i}$’s. Since $|a_{i}(0)|\leq M_{1}$ for all $i\in[m]$, we know that $|a_{i}(t)|\leq M(1+t/\varepsilon)$ for all $t\geq 0$, thus leading to the following estimate:

where the constant $M$ only depends on the $M_{i}$’s in our assumptions. At initialization, we know that $\sum_{i,j=1}^{m}r_{ij}^{\perp}(0)^{2}=m$. Applying Grönwall’s inequality yields that

To this end, we define $S(t)=\sum_{i=1}^{m}\left(\left(a_{i}(t)-a_{i}^{\mbox{\tiny\rm mf}}(t)\right)^{2}+\left(s_{i}(t)-s_{i}^{\mbox{\tiny\rm mf}}(t)\right)^{2}\right)$. By our assumption, $S(0)=0$. Moreover, using the same technique as in the proof of Lemma 1, we know that $|a_{i}^{\mbox{\tiny\rm mf}}(t)|\leq M(1+t/\varepsilon)$ for all $i\in[m]$. According to Eq.s (15)-(18) and Eq. (23), we deduce that

where in $(i)$ we use the Cauchy-Schwarz inequality and the inequality of arithmetic and geometric means, and $(ii)$ follows from the conclusion of Lemma 1. Similarly, we obtain that

Combining the above estimates, we finally deduce that

Applying Grönwall’s inequality immediately implies

which further leads to Eq. (120) and concludes the proof of Proposition 2. The “consequently” part can be shown via direct calculation, but we include it here for the sake of completeness. By definition, for any $t\in[0,T]$ we have

A.4 Derivation of the mean field dynamics (28)

where $(i)$ follows from the ODE satisfied by the $(a_{i}^{\mbox{\tiny\rm mf}}(t),s_{i}^{\mbox{\tiny\rm mf}}(t))$’s, and in $(ii)$ we use integration by parts. We thus obtain that

A.5 Details of the alternative mean field approach

where $\overline{\Psi}=(\overline{\Psi}_{a},\overline{\Psi}_{u})$ is given by

A remarkable property of the equation (124) is that it preserves invariance to rotations orthogonal to $u_{*}$. Indeed, assume that $\overline{\rho}$ is invariant to rotations orthogonal to $u_{*}$. In this case, we show that $\overline{\Psi}_{a}\left(a,u;\overline{\rho}\right)$ and $\langle u_{*},\overline{\Psi}_{u}\left(a,u;\overline{\rho}\right)\rangle$ depend only on $s:=\langle u,u_{*}\rangle$ and $s_{1}:=\langle u_{1},u_{*}\rangle$. Let $u^{\perp}$ (resp. $u_{1}^{\perp}$) denote the component of $u$ (resp. $u_{1}$) orthogonal to $u_{*}$. Let $R$ denote a random uniform rotation orthogonal to $u_{*}$. By the rotation invariance of $\overline{\rho}$,

The random variable $B^{(d)}=\left\langle\frac{u^{\perp}}{\|u^{\perp}\|},R\frac{u_{1}^{\perp}}{\|u_{1}^{\perp}\|}\right\rangle$ is a one dimensional projection of a random variable uniform on the unit sphere of the hyperplane orthogonal to $u_{*}$; thus it has the density $p_{B^{(d)}}(b)\propto(1-b^{2})^{d/2-2}$ (see, e.g., [Frye and Efthimiou, 2012, Lemma 4.17]). Denote

In the equation above, we have $\langle u_{*},(I_{d}-uu^{\top})s_{1}u_{*}\rangle=s_{1}(1-s^{2})$ and as $\langle u_{*},Ru_{1}^{\perp}\rangle=0$ a.s., we have

Of course, a discrete measure of the form (123) can not be invariant to rotations orthogonal to $u_{*}$. However, if the $u_{i}$ are initialized uniformly on the unit sphere, then the measure $\overline{\rho}_{0}$ converges to a measure with the rotation invariance as $m\to\infty$. One can then apply the results of Mei et al. to control the deviations from this limit. Let us thus assume that $\overline{\rho}_{0}$ satisfies the rotation invariance. Define the map $\varphi(a,u)=(a,\langle u,u_{*}\rangle)$. Then, from (125), (126), the push-forward $\rho_{t}$ of $\overline{\rho}_{t}$ through the map $\varphi$ satisfies the continuity equation

where $\Psi^{(d)}=(\Psi^{(d)}_{a},\Psi^{(d)}_{s})$ is given by

Appendix B Calculations for the analysis of mean-field gradient flow

In order to solve the system (83), we start from an associated one-dimensional ODE.

The solution $\lambda=\lambda(t_{3})$ of the ODE

For simplicity, denote $\alpha=|\sigma_{1}|$, $\beta=|\varphi_{1}|$ and $\gamma={|\sigma_{1}|}\left\|a_{\perp,\rm{init}}\right\|_{L^{2}(\rho)}^{2}$. Then

This is Bernoulli differential equation (see, e.g., Encyclopedia of Mathematics ). In this situation, the classical trick is to reduce the problem to a linear inhomogeneous first-order equation by considering

Solving this linear inhomogeneous first-order equation gives

Let $\lambda=\lambda(t_{3})$ be a solution of (127) and consider

Then $a^{(-1)},s^{(1)}$ are solutions of the constrained ODE system (79), (82). Indeed,

thus the constraint (79) is satisfied. Further

A similar computation shows that the differential equation for $s^{(1)}$ is also satisfied. This concludes that (129) is a valid candidate to solve the third time scale.

To determine the value of the initialization $\lambda(0)$ we perform a matching procedure with the previous time scale. In this paragraph, we denote $\underline{a},\underline{s}$ the approximation obtained in the second time scale (Section 6.2), and $\overline{a},\overline{s}$ the approximation in the third time scale (Section 6.3 and above).

Consider an intermediate time scale $\widetilde{t}=t_{2}-c\log\frac{1}{\varepsilon}$ with $0<c<\frac{1}{4|\sigma_{1}\varphi_{1}|}$. Assume $\widetilde{t}\asymp 1$. Then

From the approximation (76) on the second time scale,

From the approximation on the third time scale,

Note that as $t_{3}\to-\infty$, from (128),

By matching, Equations (130) and (131) should be coherent. This gives

One could check similarly that $s^{(1)}$ also satisfies the matching conditions under the same constraint, and thus that (129) are indeed the solutions of the third time scale.

B.2 Induced approximation of the risk

In this section, we show that the behavior of $a$ and $s$ derived in Sections 6.1–6.3 leads to an evolution of the risk alternating plateaus and rapid decreases, in agreement with the standard scenario of Definition 1. For the convenience of the reader, we recall the expression (36) of the risk

This describes, in a more detailed form, the first transition in Definition 1.

This second time scale does not induce any transition of the risk $\mathscrsfs{R}_{\mbox{\tiny\rm mf},*}$ (but was necessary to understand the divergence of $a$ and $\varepsilon^{-\nicefrac{{1}}{{2}}}s$).

where in $(a)$ we used (84) and in $(b)$ (85). Thus as $\varepsilon\to 0$,

This describes, in a more detailed form, the second transition in Definition 1.

B.3 Proof of Theorem 1

Throughout the proof, we will use the shorthand $\mathscrsfs{R}_{l}(\tau)$ to represent $\mathscrsfs{R}_{l}(\widetilde{a}(\tau),\widetilde{s}(\tau))$. First, note that according to the ODE satisfied by $\mathscrsfs{R}_{l}$ (Eq. (92)), we know that $\mathscrsfs{R}_{l}$ must be non-increasing, thus for small enough $\varepsilon>0$,

According to the comparison theorem for system of ODEs, we know that $|\widetilde{a}(\omega,\tau)|\leq\widehat{a}(\omega,\tau)$, $|\widetilde{s}(\omega,\tau)|\leq\widehat{s}(\omega,\tau)$ for all $\tau\geq 0$ where