A mean field view of the landscape of two-layer neural networks

Significance

Multilayer neural networks have proven extremely successful in a variety of tasks, from image classification to robotics. However, the reasons for this practical success and its precise domain of applicability are unknown. Learning a neural network from data requires solving a complex optimization problem with millions of variables. This is done by stochastic gradient descent (SGD) algorithms. We study the case of two-layer networks and derive a compact description of the SGD dynamics in terms of a limiting partial differential equation. Among other consequences, this shows that SGD dynamics does not become more complex when the network size increases.

Abstract

Multilayer neural networks are among the most powerful models in machine learning, yet the fundamental reasons for this success defy mathematical understanding. Learning a neural network requires optimizing a nonconvex high-dimensional objective (risk function), a problem that is usually attacked using stochastic gradient descent (SGD). Does SGD converge to a global optimum of the risk or only to a local optimum? In the former case, does this happen because local minima are absent or because SGD somehow avoids them? In the latter, why do local minima reached by SGD have good generalization properties? In this paper, we consider a simple case, namely two-layer neural networks, and prove that—in a suitable scaling limit—SGD dynamics is captured by a certain nonlinear partial differential equation (PDE) that we call distributional dynamics (DD). We then consider several specific examples and show how DD can be used to prove convergence of SGD to networks with nearly ideal generalization error. This description allows for “averaging out” some of the complexities of the landscape of neural networks and can be used to prove a general convergence result for noisy SGD.

Multilayer neural networks are one of the oldest approaches to statistical machine learning, dating back at least to the 1960s (1). Over the last 10 years, under the impulse of increasing computer power and larger data availability, they have emerged as a powerful tool for a wide variety of learning tasks (2, 3).

In this paper, we focus on the classical setting of supervised learning, whereby we are given data points $\left({\mathbf{x}}_i,y_i\right)\in {\mathbb{R}}^d\times \mathbb{R}$, indexed by $i\in \mathbb{N}$, which are assumed to be independent and identically distributed from an unknown distribution $\mathbb{P}$ on ${\mathbb{R}}^d\times \mathbb{R}$. Here ${\mathbf{x}}_i\in {\mathbb{R}}^d$ is a feature vector (e.g., a set of descriptors of an image), and $y_i\in \mathbb{R}$ is a label (e.g., labeling the object in the image). Our objective is to model the dependence of the label $y_i$ on the feature vector ${\mathbf{x}}_i$ to assign labels to previously unlabeled examples. In a two-layer neural network, this dependence is modeled as

$$ŷ\left(\mathbf{x};\mathit{\theta }\right)=\frac{1}{N}\sum _{i=1}^N{\sigma }_{*}\left(\mathbf{x};{\mathit{\theta }}_i\right).$$ [1]

Here, $N$ is the number of hidden units (neurons), ${\sigma }_{*}:{\mathbb{R}}^d\times {\mathbb{R}}^D\to \mathbb{R}$ is an activation function, and ${\mathit{\theta }}_i\in {\mathbb{R}}^D$ are parameters, which we collectively denote by $\mathit{\theta }=\left({\mathit{\theta }}_1,\dots ,{\mathit{\theta }}_N\right)$. The factor $\left(1/N\right)$ is introduced for convenience and can be eliminated by redefining the activation. Often ${\mathit{\theta }}_i=\left(a_i,b_i,{\mathbf{w}}_i\right)$ and

$${\sigma }_{*}\left(\mathbf{x};{\mathit{\theta }}_i\right)=a_i\hspace{0.17em}\sigma \left(⟨{\mathbf{w}}_i,\mathbf{x}⟩+b_i\right),$$ [2]

for some $\sigma :\mathbb{R}\to \mathbb{R}$. Ideally, the parameters $\mathit{\theta }={\left({\mathit{\theta }}_i\right)}_{i\le N}$ should be chosen as to minimize the risk (generalization error) $R_N\left(\mathit{\theta }\right)=\mathbb{E}\left\{\ell \left(y,ŷ\left(\mathbf{x};\mathit{\theta }\right)\right)\right\}$, where $\ell :\mathbb{R}\times \mathbb{R}\to \mathbb{R}$ is a certain loss function. For the sake of simplicity, we will focus on the square loss $\ell \left(y,ŷ\right)={\left(y-ŷ\right)}^2$, but more general choices can be treated along the same lines.

In practice, the parameters of neural networks are learned by stochastic gradient descent (SGD) (4) or its variants. In the present case, this amounts to the iteration

$${\mathit{\theta }}_i^{k+1}={\mathit{\theta }}_i^k+2s_k\hspace{0.17em}\left(y_k,-,ŷ,\left({\mathbf{x}}_k;{\mathit{\theta }}^k\right)\right)\hspace{0.17em}{\nabla }_{{\mathit{\theta }}_i}{\sigma }_{*}\left({\mathbf{x}}_k;{\mathit{\theta }}_i^k\right).$$ [3]

Here ${\mathit{\theta }}^k={\left({\mathit{\theta }}_i^k\right)}_{i\le N}$ denotes the parameters after $k$ iterations, $s_k$ is a step size, and $\left({\mathbf{x}}_k,y_k\right)$ is the $k$th example. Throughout the paper, we make the following “One-Pass Assumption”: Training examples are never revisited. Equivalently, ${\left\{\left({\mathbf{x}}_k,y_k\right)\right\}}_{k\ge 1}$ are independent and identically distributed. $\left({\mathbf{x}}_k,y_k\right)\sim \mathbb{P}$.

In large-scale applications, this is not far from truth: The data are so large that each example is visited at most a few times (5). Further, theoretical guarantees suggest that there is limited advantage to be gained from multiple passes (6). For recent work deriving scaling limits under such an assumption (in different problems), see ref. 7.

Understanding the optimization landscape of two-layer neural networks is largely an open problem even when we have access to an infinite number of examples—that is, to the population risk $R_N\left(\mathit{\theta }\right)$. Several studies have focused on special choices of the activation function ${\sigma }_{*}$ and of the data distribution $\mathbb{P}$, proving that the population risk has no bad local minima (8–10). This type of analysis requires delicate calculations that are somewhat sensitive to the specific choice of the model. Another line of work proposes new algorithms with theoretical guarantees (11–16), which use initializations based on tensor factorization.

In this paper, we prove that—in a suitable scaling limit—the SGD dynamics admits an asymptotic description in terms of a certain nonlinear partial differential equation (PDE). This PDE has a remarkable mathematical structure, in that it corresponds to a gradient flow in the metric space $\left(\mathcal{P}\left({\mathbb{R}}^D\right),W_2\right)$: the space of probability measures on ${\mathbb{R}}^D$, endowed with the Wasserstein metric. This gradient flow minimizes an asymptotic version of the population risk, which is defined for $\rho \in \mathcal{P}\left({\mathbb{R}}^D\right)$ and will be denoted by $R\left(\rho \right)$. This description simplifies the analysis of the landscape of two-layer neural networks, for instance by exploiting underlying symmetries. We illustrate this by obtaining results on several concrete examples as well as a general convergence result for “noisy SGD.” In the next section, we provide an informal outline, focusing on basic intuitions rather than on formal results. We then present the consequences of these ideas on a few specific examples and subsequently state our general results.

An Informal Overview

A good starting point is to rewrite the population risk $R_N\left(\mathit{\theta }\right)=\mathbb{E}\left\{{\left[y-ŷ\left(\mathbf{x};\mathit{\theta }\right)\right]}^2\right\}$ as

$$R_N\left(\mathit{\theta }\right)=R_{\#}+\frac{2}{N}\sum _{i=1}^{N}V\left({\mathit{\theta }}_i\right)+\frac{1}{N^2}\sum _{i,j=1}^{N}U\left({\mathit{\theta }}_i,{\mathit{\theta }}_j\right),$$ [4]

where we defined the potentials $V\left(\mathit{\theta }\right)=-\mathbb{E}\left\{y\hspace{0.17em}{\sigma }_{*}\left(\mathbf{x};\mathit{\theta }\right)\right\}$, $U\left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right)=\mathbb{E}\left\{{\sigma }_{*}\left(\mathbf{x};{\mathit{\theta }}_1\right){\sigma }_{*}\left(\mathbf{x};{\mathit{\theta }}_2\right)\right\}$. In particular, $U\left(\cdot ,\cdot \right)$ is a symmetric positive semidefinite kernel. The constant $R_{\#}=\mathbb{E}\left\{y^2\right\}$ is the risk of the trivial predictor $ŷ=0$.

Notice that $R_N\left(\mathit{\theta }\right)$ only depends on ${\mathit{\theta }}_1,\dots ,{\mathit{\theta }}_N$ through their empirical distribution ${\widehat{\rho }}^{\left(N\right)}=N^{-1}{\sum }_{i=1}^N{\delta }_{{\mathit{\theta }}_i}$. This suggests considering a risk function defined for $\rho \in \mathcal{P}\left({\mathbb{R}}^D\right)$ [we denote by $\mathcal{P}\left(\mathrm{\Omega }\right)$ the space of probability distributions on $\mathrm{\Omega }$]:

$$R\left(\rho \right)=R_{\#}+2\int V\left(\mathit{\theta }\right)\hspace{0.17em}\rho \left(\mathrm{\text{d}}\mathit{\theta }\right)+\int U\left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right)\hspace{0.17em}\rho \left(\mathrm{\text{d}}{\mathit{\theta }}_1\right)\hspace{0.17em}\rho \left(\mathrm{\text{d}}{\mathit{\theta }}_2\right).$$ [5]

Formal relationships can be established between $R_N\left(\mathit{\theta }\right)$ and $R\left(\rho \right)$. For instance, under mild assumptions, ${inf}_{\mathit{\theta }}{R}_N\left(\mathit{\theta }\right)={inf}_{\rho }R\left(\rho \right)+O\left(1/N\right)$. We refer to the next sections for mathematical statements of this type.

Roughly speaking, $R\left(\rho \right)$ corresponds to the population risk when the number of hidden units goes to infinity, and the empirical distribution of parameters ${\widehat{\rho }}^{\left(N\right)}$ converges to $\rho$. Since $U\left(\cdot ,\cdot \right)$ is positive semidefinite, we obtain that the risk becomes convex in this limit. The fact that learning can be viewed as convex optimization in an infinite-dimensional space was indeed pointed out in the past (17, 18). Does this mean that the landscape of the population risk simplifies for large $N$ and descent algorithms will converge to a unique (or nearly unique) global optimum?

The answer to the latter question is generally negative, and a physics analogy can explain why. Think of ${\mathit{\theta }}_1,\dots ,{\mathit{\theta }}_N$ as the positions of $N$ particles in a $D$-dimensional space. When $N$ is large, the behavior of such a “gas” of particles is effectively described by a density ${\rho }_t\left(\mathit{\theta }\right)$ (with $t$ indexing time). However, not all “small” changes of this density profile can be realized in the actual physical dynamics: The dynamics conserves mass locally because particles cannot move discontinuously. For instance, if $\mathrm{\text{supp}}\left({\rho }_t\right)=S_1\cup S_2$ for two disjoint compact sets $S_1,S_2\subseteq {\mathbb{R}}^D$ and all $t\in \left[t_1,t_2\right]$, then the total mass in each of these regions cannot change over time—that is, ${\rho }_t\left(S_1\right)=1-{\rho }_t\left(S_2\right)$ does not depend on $t\in \left[t_1,t_2\right]$.

We will prove that SGD is well approximated (in a precise quantitative sense described below) by a continuum dynamics that enforces this local mass conservation principle. Namely, assume that the step size in SGD is given by $s_k=\epsilon \xi \left(k\epsilon \right)$, for $\xi :{\mathbb{R}}_{\ge 0}\to {\mathbb{R}}_{\ge 0}$ a sufficiently regular function. Denoting by ${\widehat{\rho }}_k^{\left(N\right)}=N^{-1}{\sum }_{i=1}^N{\delta }_{{\mathit{\theta }}_i^k}$ the empirical distribution of parameters after $k$ SGD steps, we prove that

$${\widehat{\rho }}_{t/\epsilon }^{\left(N\right)}\Rightarrow {\rho }_t,$$ [6]

when $N\to \infty$, $\epsilon \to 0$ (here $\Rightarrow$ denotes weak convergence). The asymptotic dynamics of ${\rho }_t$ is defined by the following PDE, which we shall refer to as distributional dynamics (DD):

$${\partial }_t{\rho }_t=2\xi \left(t\right)\hspace{0.17em}{\nabla }_{\mathit{\theta }}\cdot \left({\rho }_t{\nabla }_{\mathit{\theta }}\mathrm{\Psi }\left(\mathit{\theta };{\rho }_t\right)\right),$$ [7]

and

$$\mathrm{\Psi }\left(\mathit{\theta };\rho \right)\equiv V\left(\mathit{\theta }\right)+\int U\left(\mathit{\theta },\mathit{\theta }\prime \right)\rho \left(\mathrm{\text{d}}\mathit{\theta }\prime \right).$$ [8]

[Here, ${\nabla }_{\mathit{\theta }}\cdot \mathbf{v}\left(\mathit{\theta }\right)$ denotes the divergence of the vector field $\mathbf{v}\left(\mathit{\theta }\right)$.] This should be interpreted as an evolution equation in $\mathcal{P}\left({\mathbb{R}}^D\right)$. While we described the convergence to this dynamics in asymptotic terms, the results in the next sections provide explicit nonasymptotic bounds. In particular, ${\rho }_t$ is a good approximation of ${\widehat{\rho }}_k^{\left(N\right)}$, $k=t/\epsilon$, as soon as $\epsilon \ll 1/D$ and $N\gg D$.

Using these results, analyzing learning in two-layer neural networks reduces to analyzing the PDE (Eq. 7). While this is far from being an easy task, the PDE formulation leads to several simplifications and insights. First, it factors out the invariance of the risk (Eq. 4) (and of the SGD dynamics; Eq. 3), with respect to permutations of the units $\left\{1,\dots ,N\right\}$.

Second, it allows us to exploit symmetries in the data distribution $\mathbb{P}$. If $\mathbb{P}$ is left invariant under a group of transformations (e.g., rotations), we can look for a solution ${\rho }_t$ of the DD (Eq. 7) that enjoys the same symmetry, hence reducing the dimensionality of the problem. This is impossible for the finite-$N$ dynamics (Eq. 3), since no arrangement of the points $\left\{{\mathit{\theta }}_1,\dots ,{\mathit{\theta }}_N\right\}\subseteq {\mathbb{R}}^D$ is left invariant, say, under rotations. We will provide examples of this approach in the next sections.

Third, there is rich mathematical literature on the PDE (Eq. 7) that was motivated by the study of interacting particle systems in mathematical physics. As mentioned above, a key structure exploited in this line of work is that Eq. 7 can be viewed as a gradient flow for the cost function $R\left(\rho \right)$ in the space $\left(\mathcal{P}\left({\mathbb{R}}^D\right),W_2\right)$, of probability measures on ${\mathbb{R}}^D$ endowed with the Wasserstein metric (19–21). Roughly speaking, this means that the trajectory $t\mapsto {\rho }_t$ attempts to minimize the risk $R\left(\rho \right)$ while maintaining the “local mass conservation” constraint. Recall that the Wasserstein distance is defined as

$$W_2\left({\rho }_1,{\rho }_2\right)={\left(\underset{\gamma \in \mathcal{C}\left({\rho }_1,{\rho }_2\right)}{inf}\int \Vert {\mathit{\theta }}_1-{\mathit{\theta }}_2{\Vert }_2^2\gamma \left(\mathrm{\text{d}}{\mathit{\theta }}_1,\mathrm{\text{d}}{\mathit{\theta }}_2\right)\hspace{0.17em}\right)}^{1/2},$$ [9]

where the infimum is taken over all couplings of ${\rho }_1$ and ${\rho }_2$. Informally, the fact that ${\rho }_t$ is a gradient flow means that Eq. 7 is equivalent, for small $\tau$, to

$${\rho }_{t+\tau }\approx \arg \underset{\rho \in \mathcal{P}\left({\mathbb{R}}^D\right)}{min}\left\{R\left(\rho \right)+\frac{1}{2\xi \left(t\right)\tau }{W}_2{\left(\rho ,{\rho }_t\right)}^2\right\}.$$ [10]

Powerful tools from the mathematical literature on gradient flows in measure spaces (20) can be exploited to study the behavior of Eq. 7.

Most importantly, the scaling limit elucidates the dependence of the landscape of two-layer neural networks on the number of hidden units $N$.

A remarkable feature of neural networks is the observation that, while they might be dramatically overparametrized, this does not lead to performance degradation. In the case of bounded activation functions, this phenomenon was clarified in the 1990s for empirical risk minimization algorithms (see, e.g., ref. 22). The present work provides analogous insight for the SGD dynamics: Roughly speaking, our results imply that the landscape remains essentially unchanged as $N$ grows, provided $N\gg D$. In particular, assume that the PDE (Eq. 7) converges close to an optimum in time $t_{*}\left(D\right)$. This might depend on $D$ but does not depend on the number of hidden units $N$ (which does not appear in the DD PDE; Eq. 7). If $t_{*}\left(D\right)=O_D\left(1\right)$, we can then take $N$ arbitrarily (as long as $N\gg D$) and will achieve a population risk that is independent of $N$ (and corresponds to the optimum), using $k=t_{*}/\epsilon =O\left(D\right)$ samples.

Our analysis can accommodate some important variants of SGD, a particularly interesting one being noisy SGD:

$$\begin{array}{ccc}\hfill {\mathit{\theta }}_i^{k+1}& =\hfill & \left(1-2\lambda s_k\right){\mathit{\theta }}_i^k+2s_k\hspace{0.17em}\left(y_k-{ŷ}_k\right)\hspace{0.17em}{\nabla }_{{\mathit{\theta }}_i}{\sigma }_{*}\left({\mathbf{x}}_k;{\mathit{\theta }}_i^k\right)\hfill \\ \hfill & \hfill & +\hspace{0.17em}\sqrt{2s_k/\beta }\hspace{0.17em}{\mathbf{g}}_i^k,\hfill \end{array}$$ [11]

where ${\mathbf{g}}_i^k\sim \mathsf{\text{N}}\left(0,{\mathbf{I}}_D\right)$ and ${ŷ}_k=ŷ\left({\mathbf{x}}_k;{\mathit{\theta }}^k\right)$. (The term $-2\lambda s_k{\mathit{\theta }}_i^k$ corresponds to an ${\ell }_2$ regularization and will be useful for our analysis below.) The resulting scaling limit differs from Eq. 7 by the addition of a diffusion term:

$${\partial }_t{\rho }_t=2\xi \left(t\right)\hspace{0.17em}{\nabla }_{\mathit{\theta }}\cdot \left({\rho }_t{\nabla }_{\mathit{\theta }}{\mathrm{\Psi }}_{\lambda }\left(\mathit{\theta };{\rho }_t\right)\right)+2\xi \left(t\right){\beta }^{-1}{\mathrm{\Delta }}_{\mathit{\theta }}{\rho }_t,$$ [12]

where ${\mathrm{\Psi }}_{\lambda }\left(\mathit{\theta };\rho \right)=\mathrm{\Psi }\left(\mathit{\theta };\rho \right)+\left(\lambda /2\right)\Vert \mathit{\theta }{\Vert }_2^2$, and ${\mathrm{\Delta }}_{\mathit{\theta }}f\left(\mathit{\theta }\right)={\sum }_{i=1}^d{\partial }_{{\theta }_i}^{2}f\left(\mathit{\theta }\right)$ denotes the usual Laplacian. This can be viewed as a gradient flow for the free-energy $F_{\beta ,\lambda }\left(\rho \right)=\left(1/2\right)R\left(\rho \right)+\left(\lambda /2\right)\int \Vert \mathit{\theta }{\Vert }_2^2\rho \left(\mathrm{\text{d}}\mathit{\theta }\right)-{\beta }^{-1}\mathrm{\text{Ent}}\left(\rho \right)$, where $\mathrm{\text{Ent}}\left(\rho \right)=-\int \rho \left(\mathit{\theta }\right)\log \rho \left(\mathit{\theta }\right)\hspace{0.17em}\mathrm{\text{d}}\mathit{\theta }$ is the entropy of $\rho$ [by definition $\mathrm{\text{Ent}}\left(\rho \right)=-\infty$ if $\rho$ is singular]. $F_{\beta ,\lambda }\left(\rho \right)$ is an entropy-regularized risk, which penalizes strongly nonuniform $\rho$.

We will prove below that, for $\beta <\infty$, the evolution (Eq. 12) generically converges to the minimizer of $F_{\beta ,\lambda }\left(\rho \right)$, hence implying global convergence of noisy SGD in a number of steps independent of $N$.

Examples

In this section, we discuss some simple applications of the general approach outlined above. Let us emphasize that these examples are not realistic. First, the data distribution $\mathbb{P}$ is extremely simple: We made this choice to be able to carry out explicit calculations. Second, the activation function ${\sigma }_{*}\left(\mathbf{x};\mathit{\theta }\right)$ is not necessarily optimal: We made this choice to illustrate some interesting phenomena.

Centered Isotropic Gaussians.

One-neuron neural networks perform well with (nearly) linearly separable data. The simplest classification problem that requires multilayer networks is, arguably, the one of distinguishing two Gaussians with the same mean. Assume the joint law $\mathbb{P}$ of $\left(y,\mathbf{x}\right)$ to be as follows:

with probability $1/2$: $y=+1$, $\mathbf{x}\sim \mathsf{\text{N}}\left(0,{\left(1+\mathrm{\Delta }\right)}^2{\mathbf{I}}_d\right)$; and

with probability $1/2$: $y=-1$, $\mathbf{x}\sim \mathsf{\text{N}}\left(0,{\left(1-\mathrm{\Delta }\right)}^2{\mathbf{I}}_d\right)$.

(This example will be generalized later.) Of course, optimal classification in this model becomes entirely trivial if we compute the feature $h\left(\mathbf{x}\right)=\Vert \mathbf{x}{\Vert }_2$. However, it is nontrivial that an SGD-trained neural network will succeed.

We choose an activation function without offset or output weights, namely ${\sigma }_{*}\left(\mathbf{x};{\mathit{\theta }}_i\right)=\sigma \left(⟨{\mathbf{w}}_i,\mathbf{x}⟩\right)$. While qualitatively similar results are obtained for other choices of $\sigma$, we will use a simple piecewise linear function as a running example: $\sigma \left(t\right)=s_1$ if $t\le t_1$, $\sigma \left(t\right)=s_2$ if $t\ge t_2$, and $\sigma \left(t\right)$ interpolated linearly for $t\in \left(t_1,t_2\right)$. In simulations, we use $t_1=0.5$, $t_2=1.5$, $s_1=-2.5$, and $s_2=7.5$.

We run SGD with initial weights ${\left({\mathbf{w}}_i^0\right)}_{i\le N}{\sim }_{iid}{\rho }_0$, where ${\rho }_0$ is spherically symmetric. Fig. 1 reports the result of such an experiment. Due to the symmetry of the distribution $\mathbb{P}$, the distribution ${\rho }_t$ remains spherically symmetric for all $t$ and hence is completely determined by the distribution ${\overline{\rho }}_t$ of the norm $r=\Vert \mathbf{w}{\Vert }_2$. This distribution satisfies a one-dimensional reduced DD:

$${\partial }_t{\overline{\rho }}_t=2\xi \left(t\right)\hspace{0.17em}{\partial }_r\left({\overline{\rho }}_t{\partial }_r\psi \left(r;{\overline{\rho }}_t\right)\right),$$ [13]

where the form of $\psi \left(r;\rho \right)$ can be derived from $\mathrm{\Psi }\left(\mathit{\theta };\rho \right)$. This reduced PDE can be efficiently solved numerically, see SI Appendix for technical details. As illustrated by Fig. 1, the empirical results match closely the predictions produced by this PDE.

Fig. 1.

Evolution of the radial distribution ${\overline{\rho }}_t$ for the isotropic Gaussian model, with $\mathrm{\Delta }=0.8$. Histograms are obtained from SGD experiments with $d=40$, $N=800$, initial weight distribution ${\rho }_0=\mathsf{\text{N}}\left(\mathbf{0},0.8^2/d\cdot {\mathbf{I}}_d\right)$, and step size $ϵ=10^{-6}$ and $\xi \left(t\right)=1$. Continuous lines correspond to a numerical solution of the DD (Eq. 13).

In Fig. 2, we compare the asymptotic risk achieved by SGD with the prediction obtained by minimizing $R\left(\rho \right)$ (cf., Eq. 5) over spherically symmetric distributions. It turns out that, for certain values of $\mathrm{\Delta }$, the minimum is achieved by the uniform distribution over a sphere of radius $\Vert \mathbf{w}{\Vert }_2=r_{*}$, to be denoted by ${\rho }_{r_{*}}^{\mathrm{\text{unif}}}$. The value of $r_{*}$ is computed by minimizing

$${\overline{R}}_d^{\left(1\right)}\left(r\right)=1+2v\left(r\right)\hspace{0.17em}+u_d\left(r,r\right),$$ [14]

where expressions for $v\left(r\right)$, $u_d\left(r_1,r_2\right)$ can be readily derived from $V\left(\mathbf{w}\right)$, $U\left({\mathbf{w}}_1,{\mathbf{w}}_2\right)$ and are given in SI Appendix.

Fig. 2.

Population risk in the problem of separating two isotropic Gaussians, as a function of the separation parameter $\mathrm{\Delta }$. We use a two-layer network with piecewise linear activation, no offset, and output weights equal to 1. Empirical results obtained by SGD (a single run per data point) are marked “+.” Continuous lines are theoretical predictions obtained by numerically minimizing $R\left(\rho \right)$ (see SI Appendix for details). Dashed lines are theoretical predictions from the single-delta ansatz of Eq. 14. Notice that this ansatz is incorrect for $\mathrm{\Delta }>{\mathrm{\Delta }}_d^{\mathrm{h}}$, which is marked as a solid round dot. Here, $N=800$.

Lemma 1:

Let $r_{*}$ be a global minimizer of $r\mapsto R_d^{\left(1\right)}\left(r\right)$. Then ${\rho }_{r_{*}}^{\mathrm{\text{unif}}}$ is a global minimizer of $\rho \mapsto R\left(\rho \right)$ if and only if $v\left(r\right)+u_d\left(r,r_{*}\right)\ge v\left(r_{*}\right)+u_d\left(r_{*},r_{*}\right)$ for all $r\ge 0$.

Checking numerically, this condition yields that ${\rho }_{r_{*}}^{\mathrm{\text{unif}}}$ is a global minimizer for $\mathrm{\Delta }$ in an interval $\left[{\mathrm{\Delta }}_d^{\mathrm{l}},{\mathrm{\Delta }}_d^{\mathrm{h}}\right]$, where ${lim}_{d\to \infty }{\mathrm{\Delta }}_d^{\mathrm{l}}=0$ and ${lim}_{d\to \infty }{\mathrm{\Delta }}_d^{\mathrm{h}}={\mathrm{\Delta }}_{\infty }\approx 0.47$.

Fig. 2 shows good quantitative agreement between empirical results and theoretical predictions and suggests that SGD achieves a value of the risk that is close to optimum. Can we prove that this is indeed the case and that the SGD dynamics does not get stuck in local minima? It turns out that we can use our general theory (see next section) to prove that this is the case for large $d$. To state this result, we need to introduce a class of good uninformative initializations ${\mathcal{P}}_{\mathrm{\text{good}}}\subseteq \mathcal{P}\left({\mathbb{R}}_{\ge 0}\right)$ for which convergence to the optimum takes place. For $\overline{\rho }\in \mathcal{P}\left({\mathbb{R}}_{\ge 0}\right)$, we let ${\overline{R}}_d\left(\overline{\rho }\right)\equiv R\left(\overline{\rho }\times \mathrm{U}\mathrm{n}\mathrm{i}\mathrm{f}\left({\mathbb{S}}^{d-1}\right)\right)$. This risk has a well-defined limit as $d\to \infty$. We say that $\overline{\rho }\in {\mathcal{P}}_{\mathrm{\text{good}}}$ if $\left(i\right)\hspace{0.17em}\overline{\rho }$ is absolutely continuous with respect to the Lebesgue measure, with bounded density, $\left(ii\right)\hspace{0.17em}{\overline{R}}_{\infty }\left(\overline{\rho }\right)<1$.

Theorem 1:

For any $\eta ,\mathrm{\Delta },\delta >0$ and ${\overline{\rho }}_0\in {\mathcal{P}}_{\mathrm{\text{good}}}$, there exists $d_0=d_0\left(\eta ,{\overline{\rho }}_0,\mathrm{\Delta }\right)$, $T=T\left(\eta ,{\overline{\rho }}_0,\mathrm{\Delta }\right)$, and $C_0=C_0\left(\eta ,{\overline{\rho }}_0,\mathrm{\Delta },\delta \right)$, such that the following holds for the problem of classifying isotropic Gaussians. For any dimension $d\ge d_0$, number of neurons $N\ge C_{0}d$, consider SGD initialized with ${\left({\mathbf{w}}_i^0\right)}_{i\le N}{\sim }_{iid}{\overline{\rho }}_0\times \mathrm{U}\mathrm{n}\mathrm{i}\mathrm{f}\left({\mathbb{S}}^{d-1}\right)$ and step size $\epsilon \in \left[1/N^{10},1/\left(C_{0}d\right)\right]$. Then we have

$$R_N\left({\mathit{\theta }}^k\right)\le \underset{\mathit{\theta }\in {\mathbb{R}}^{N\times d}}{inf}R\left(\mathit{\theta }\right)+\eta$$ [15]

for any $k\in \left[T/\epsilon ,10T/\epsilon \right]$ with probability at least $1-\delta$.

In particular, if we set $\epsilon =1/\left(C_{0}d\right)$, then the number of SGD steps is $k\in \left[\left(C_{0}T\right)\hspace{0.17em}d,\left(10C_{0}T\right)\hspace{0.17em}d\right]$: The number of samples used by SGD does not depend on the number of hidden units $N$ and is only linear in the dimension. Unfortunately, the proof does not provide the dependence of $T$ on $\eta$, but Theorem 6 below suggests exponential local convergence.

While we stated Theorem 1 for the piecewise linear sigmoids, SI Appendix presents technical conditions under which it holds for a general monotone function $\sigma :\mathbb{R}\to \mathbb{R}$.

Centered Anisotropic Gaussians.

We can generalize the previous result to a problem in which the network needs to select a subset of relevant nonlinear features out of many a priori equivalent ones. We assume the joint law of $\left(y,\mathbf{x}\right)$ to be as follows:

with probability $1/2$: $y=+1$, $\mathbf{x}\sim \mathsf{\text{N}}\left(0,{\mathit{\Sigma }}_{+}\right)$; and

with probability $1/2$: $y=-1$, $\mathbf{x}\sim \mathsf{\text{N}}\left(0,{\mathit{\Sigma }}_{-}\right)$.

Given a linear subspace $\mathcal{V}\subseteq {\mathbb{R}}^d$ of dimension $s_0\le d$, we assume that ${\mathit{\Sigma }}_{+}$, ${\mathit{\Sigma }}_{-}$ differ uniquely along $\mathcal{V}$: ${\mathit{\Sigma }}_{\pm }={\mathbf{I}}_d+\left({\tau }_{\pm }^2-1\right){\mathbf{P}}_{\mathcal{V}}$, where ${\tau }_{\pm }=\left(1\pm \mathrm{\Delta }\right)$ and ${\mathbf{P}}_{\mathcal{V}}$ is the orthogonal projector onto $\mathcal{V}$. In other words, the projection of $\mathbf{x}$ on the subspace $\mathcal{V}$ is distributed according to an isotropic Gaussian with variance ${\tau }_{+}^2$ (if $y=+1$) or ${\tau }_{-}^2$ (if $y=-1$). The projection orthogonal to $\mathcal{V}$ has instead the same variance in the two classes. A successful classifier must be able to learn the relevant subspace $\mathcal{V}$. We assume the same class of activations ${\sigma }_{*}\left(\mathbf{x};\mathit{\theta }\right)=\sigma \left(⟨\mathbf{w},\mathbf{x}⟩\right)$ as for the isotropic case.

The distribution $\mathbb{P}$ is invariant under a reduced symmetry group $\mathcal{O}\left(s_0\right)\times \mathcal{O}\left(d-s_0\right)$. As a consequence, letting $r_1=\Vert {\mathbf{P}}_{\mathcal{V}}\mathbf{w}{\Vert }_2$ and $r_2\equiv \Vert \left({\mathbf{I}}_d-{\mathbf{P}}_{\mathcal{V}}\right)\mathbf{w}{\Vert }_2$, it is sufficient to consider distributions $\rho$ that are uniform, conditional on the values of $r_1$ and $r_2$. If we initialize ${\rho }_0$ to be uniform conditional on $\left(r_1,r_2\right)$, this property is preserved by the evolution (Eq. 7). As in the isotropic case, we can use our general theory to prove convergence to a near-optimum if $d$ is large enough.

Theorem 2:

For any $\eta ,\mathrm{\Delta },\delta >0$ and ${\overline{\rho }}_0\in {\mathcal{P}}_{\mathrm{\text{good}}}$, there exists $d_0=d_0\left(\eta ,{\overline{\rho }}_0,\mathrm{\Delta },\gamma \right)$, $T=T\left(\eta ,{\overline{\rho }}_0,\mathrm{\Delta },\gamma \right)$, and $C_0=C_0\left(\eta ,{\overline{\rho }}_0,\mathrm{\Delta },\delta ,\gamma \right)$, such that the following holds for the problem of classifying anisotropic Gaussians with $s_0=\gamma d$, $\gamma \in \left(\mathrm{0,1}\right)$ fixed. For any dimension parameters $s_0=\gamma d\ge d_0$, number of neurons $N\ge C_{0}d$, consider SGD initialized with initialization ${\left({\mathbf{w}}_i^0\right)}_{i\le N}{\sim }_{iid}{\overline{\rho }}_0\times \mathrm{U}\mathrm{n}\mathrm{i}\mathrm{f}\left({\mathbb{S}}^{d-1}\right)$ and step size $\epsilon \in \left[1/N^{10},1/\left(C_{0}d\right)\right]$. Then, we have $R_N\left({\mathit{\theta }}^k\right)\le \underset{\mathit{\theta }\in {\mathbb{R}}^{N\times d}}{inf}{R}_N\left(\mathit{\theta }\right)+\eta$ for any $k\in \left[T/\epsilon ,10T/\epsilon \right]$ with probability at least $1-\delta$.

Even with a reduced degree of symmetry, SGD converges to a network with nearly optimal risk, after using a number of samples $k=O\left(d\right)$, which is independent of the number of hidden units $N$.

A Better Activation Function.

Our previous examples use activation functions ${\sigma }_{*}\left(\mathbf{x};\mathit{\theta }\right)=\sigma \left(⟨\mathbf{w},\mathbf{x}⟩\right)$ without output weights or offset to simplify the analysis and illustrate some interesting phenomena. Here we consider instead a standard rectified linear unit (ReLU) activation and fit both the output weight and the offset: ${\sigma }_{*}\left(\mathbf{x};\mathit{\theta }\right)=a\hspace{0.17em}{\sigma }_{\mathrm{\text{ReLU}}}\left(⟨\mathbf{w},\mathbf{x}⟩+b\right)$, where ${\sigma }_{\mathrm{\text{ReLU}}}\left(x\right)=max\left(x,0\right)$. Hence, $\mathit{\theta }=\left(\mathbf{w},a,b\right)\in {\mathbb{R}}^{d+2}$.

We consider the same data distribution introduced in the last section (anisotropic Gaussians). Fig. 3 reports the evolution of the risk $R_N\left({\mathit{\theta }}^k\right)$ for three experiments with $d=320$, $s_0=60$, and different values of $\mathrm{\Delta }$. SGD is initialized by setting $a_i=1$, $b_i=1$, and ${\mathbf{w}}_i^0{\sim }_{iid}\mathsf{\text{N}}\left(\mathbf{0},0.8^2/d\cdot {\mathbf{I}}_d\right)$ for $i\le N$. We observe that SGD converges to a network with very small risk, but this convergence has a nontrivial structure and presents long flat regions.

Fig. 3.

Evolution of the population risk for the variable selection problem using a two-layer neural network with ReLU activations. Here $d=320$, $s_0=60$, and $N=800$, and we used $\xi \left(t\right)=t^{-1/4}$ and $\epsilon =2\times 10^{-4}$ to set the step size. Numerical simulations using SGD (one run per data point) are marked +, and curves are solutions of the reduced PDE with $d=\infty$. (Inset) Evolution of three parameters of the reduced distribution ${\overline{\rho }}_t$ (average output weights $a$, average offsets $b$, and average ${\ell }_2$ norm in the relevant subspace $r_1$) for the same setting.

The empirical results are well captured by our predictions based on the continuum limit. In this case, we obtain a reduced PDE for the joint distribution of the four quantities $\mathbf{r}=\left(a,b,r_1=\Vert {\mathbf{P}}_{\mathcal{V}}\mathbf{w}{\Vert }_2,r_2=\Vert {\mathbf{P}}_{\mathcal{V}}^{\perp }{\mathbf{w}\Vert }_2\right)$, denoted by ${\overline{\rho }}_t$. The reduced PDE is analogous to Eq. 13 albeit in 4 dimensions rather than 1 dimension. In Fig. 3, we consider the evolution of the risk, alongside three properties of the distribution ${\overline{\rho }}_t$—the means of the output weight $a$, of the offset $b$, and of $r_1$.

Predicting Failure.

SGD does not always converge to a near global optimum. Our analysis allows us to construct examples in which SGD fails. For instance, Fig. 4 reports results for the isotropic Gaussians problem. We violate the assumptions of Theorem 1 by using nonmonotone activation function. Namely, we use ${\sigma }_{*}\left(\mathbf{x};\mathit{\theta }\right)=\sigma \left(⟨\mathbf{w},\mathbf{x}⟩\right)$, where $\sigma \left(t\right)=-2.5$ for $t\le 0$, $\sigma \left(t\right)=7.5$ for $t\ge 1.5$, and $\sigma \left(t\right)$ linearly interpolates from (0, –2.5) to (0.5, –4), and from (0.5, –4) to $\left(1.5,7.5\right)$.

Fig. 4.

Separating two isotropic Gaussians, with a nonmonotone activation function (see Predicting Failure for details). Here $N=800$, $d=320$, and $\mathrm{\Delta }=0.5$. The main frame presents the evolution of the population risk along the SGD trajectory, starting from two different initializations of ${\left({\mathbf{w}}_i^0\right)}_{i\le N}{\sim }_{iid}\mathsf{N}\left(\mathbf{0},{\kappa }^2/d\cdot {\mathbf{I}}_d\right)$ for either $\kappa =0.1$ or $\kappa =0.4$. In Inset, we plot the evolution of the average of $\Vert \mathbf{w}{\Vert }_2$ for the same conditions. Symbols are empirical results. Continuous lines are predictions obtained with the reduced PDE (Eq. 13).

Depending on the initialization, SGD converges to two different limits, one with a small risk and the second with high risk. Again, this behavior is well tracked by solving a one-dimensional PDE for the distribution ${\overline{\rho }}_t$ of $r=\Vert \mathbf{w}{\Vert }_2$.

General Results

In this section, we return to the general supervised learning problem described in the Introduction and describe our general results. Proofs are deferred to SI Appendix.

First, we note that the minimum of the asymptotic risk $R\left(\rho \right)$ of Eq. 5 provides a good approximation of the minimum of the finite-$N$ risk $R_N\left(\mathit{\theta }\right)$.

Proposition 1:

Assume that either one of the following conditions hold: $\left(a\right)\hspace{0.17em}{inf}_{\rho }R\left(\rho \right)$ is achieved by a distribution ${\rho }_{*}$ such that $\int U\left(\mathit{\theta },\mathit{\theta }\right)\hspace{0.17em}{\rho }_{*}\left(\mathrm{\text{d}}\mathit{\theta }\right)\le K$; $\left(b\right)$ There exists ${\epsilon }_0>0$ such that, for any $\rho \in \mathcal{P}\left({\mathbb{R}}^D\right)$ such that $R\left(\rho \right)\le {inf}_{\rho }R\left(\rho \right)+{\epsilon }_0$, we have $\int U\left(\mathit{\theta },\mathit{\theta }\right)\hspace{0.17em}\rho \left(\mathrm{\text{d}}\mathit{\theta }\right)\le K$. Then

$$\left|\underset{\mathit{\theta }}{inf},R_N,\left(\mathit{\theta }\right),-,\underset{\rho }{inf},R,\left(\rho \right)\right|\le K/N.$$ [16]

Further, assume that $\mathit{\theta }\mapsto V\left(\mathit{\theta }\right)$ and $\left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right)\mapsto U\left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right)$ are continuous, with $U$ bounded below. A probability measure ${\rho }_{*}$ is a global minimum of $R$ if ${inf}_{\mathit{\theta }\in {\mathbb{R}}^D}\mathrm{\Psi }\left(\mathit{\theta };{\rho }_{*}\right)>-\infty$ and

$$\mathrm{\text{supp}}\left({\rho }_{*}\right)\subseteq \arg \underset{\mathit{\theta }\in {\mathbb{R}}^D}{min}\mathrm{\Psi }\left(\mathit{\theta };{\rho }_{*}\right).$$ [17]

We next consider the DDs (Eqs. 7 and 12). These should be interpreted to hold in a weak sense (cf. SI Appendix). To establish that these PDEs indeed describe the limit of the SGD dynamics, we make the following assumptions:

• A1. $t\mapsto \xi \left(t\right)$ is bounded Lipschitz: $\Vert \xi {\Vert }_{\infty },\Vert \xi {\Vert }_{\mathrm{\text{Lip}}}\le K_1$, with ${\int }_0^{\infty }\xi \left(t\right)\mathrm{\text{d}}t=\infty$.

• A2. The activation function $\left(\mathbf{x},\mathit{\theta }\right)\mapsto {\sigma }_{*}\left(\mathbf{x};\mathit{\theta }\right)$ is bounded, with sub-Gaussian gradient: $\Vert {\sigma }_{*}{\Vert }_{\infty }\le K_2$, $\Vert {\nabla }_{\mathit{\theta }}{\sigma }_{*}\left(\mathbf{X};\mathit{\theta }\right){\Vert }_{{\psi }_2}\le K_2$. Labels are bounded $\left|y_k\right|\le K_2$.

• A3. The gradients $\mathit{\theta }\mapsto \nabla V\left(\mathit{\theta }\right)$, $\left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right)\mapsto {\nabla }_{{\mathit{\theta }}_1}U\left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right)$ are bounded, Lipschitz continuous [namely $\Vert {\nabla }_{\mathit{\theta }}V\left(\mathit{\theta }\right){\Vert }_2$, $\Vert {\nabla }_{{\mathit{\theta }}_1}U\left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right){\Vert }_2\le K_3$, $\Vert {\nabla }_{\mathit{\theta }}V\left(\mathit{\theta }\right)-{\nabla }_{\mathit{\theta }}V\left(\mathit{\theta }\prime \right){\Vert }_2\le K_3\Vert \mathit{\theta }-\mathit{\theta }\prime {\Vert }_2$, $\Vert {\nabla }_{{\mathit{\theta }}_1}U\left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right)-{\nabla }_{{\mathit{\theta }}_1}U\left({\mathit{\theta }}_1^{\prime },{\mathit{\theta }}_2^{\prime }\right){\Vert }_2\le K_3\Vert \left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right)-\left({\mathit{\theta }}_1^{\prime },{\mathit{\theta }}_2^{\prime }\right){\Vert }_2$].

We also introduce the following error term that quantifies in a nonasymptotic sense the accuracy of our PDE model:

$${\mathsf{\text{err}}}_{N,D}\left(z\right)\equiv \sqrt{1/N\vee \epsilon }\cdot \left[\sqrt{D+\log \left(N/\epsilon \right)}+z\right].$$ [18]

The convergence of the SGD process to the PDE model is an example of a phenomenon that is known in probability theory as propagation of chaos (23).

Theorem 3:

Assume that conditions A1, A2, A3 hold. For ${\rho }_0\in \mathcal{P}\left({\mathbb{R}}^D\right)$, consider SGD with initialization ${\left({\mathit{\theta }}_i^0\right)}_{i\le N}{\sim }_{iid}{\rho }_0$ and step size $s_k=\epsilon \xi \left(k\epsilon \right)$. For $t\ge 0$, let ${\rho }_t$ be the solution of PDE (Eq. 7). Then, for any fixed $k$, ${\widehat{\rho }}_k^{\left(N\right)}\Rightarrow {\rho }_{k\epsilon }$ almost surely along any sequence $\left(N,\epsilon ={\epsilon }_N\right)$ such that $N/\log \left(1/{\epsilon }_N\right)\to \infty$, ${\epsilon }_N\to 0$. Further, there exists a constant $C$ (depending uniquely on the parameters $K_i$ of conditions A1–A3) such that, for any $f:{\mathbb{R}}^D\times \mathbb{R}\to \mathbb{R}$, with $\Vert f{\Vert }_{\infty },\Vert f{\Vert }_{\mathrm{\text{Lip}}}\le 1$, $\epsilon \le 1$,

$$\begin{array}{ll}\hfill & \underset{k\in \left[0,T/\epsilon \right]\cap \mathbb{N}}{sup}\left|\frac{1}{N}\sum _{i=1}^{N}f\left({\mathit{\theta }}_i^k\right)-\int \hspace{0.17em}f\left(\mathit{\theta }\right){\rho }_{k\epsilon }\left(\mathrm{\text{d}}\mathit{\theta }\right)\right|\le C{e}^{CT}\hspace{0.17em}{\mathsf{\text{err}}}_{N,D}\left(z\right),\hfill \\ \hfill & \underset{k\in \left[0,T/\epsilon \right]\cap \mathbb{N}}{sup}\left|R_N\left({\mathit{\theta }}^k\right)-R\left({\rho }_{k\epsilon }\right)\right|\le C{e}^{CT}\hspace{0.17em}{\mathsf{\text{err}}}_{N,D}\left(z\right),\hfill \end{array}$$ [19]

with probability $1-e^{-z^2}$. The same statements hold for noisy SGD (Eq. 11), provided Eq. 7 is replaced by Eq. 12, and if $\beta \ge 1$, $\lambda \le 1$, and ${\rho }_0$ is $K_0$ sub-Gaussian for some $K_0>0$.

Notice that dependence of the error terms in $N$ and $D$ is rather benign. On the other hand, the error grows exponentially with the time horizon $T$, which limits its applicability to cases in which the DD converges rapidly to a good solution. We do not expect this behavior to be improvable within the general setting of 0.3, which a priori includes cases in which the dynamics is unstable.

We can regard $\mathbf{J}\left(\mathit{\theta };{\rho }_t\right)={\rho }_t\left(\mathit{\theta }\right){\nabla }_{\mathit{\theta }}\mathrm{\Psi }\left(\mathit{\theta };{\rho }_t\right)$ as a current. The fixed points of the continuum dynamics are densities that correspond to zero current, as stated below.

Proposition 2:

Assume $V\left(\cdot \right),U\left(\cdot ,\cdot \right)$ to be differentiable with bounded gradient. If ${\rho }_t$ is a solution of the PDE (Eq. 7), then $R\left({\rho }_t\right)$ is nonincreasing. Further, probability distribution $\rho$ is a fixed point of the PDE (Eq. 7) if and only if

$$\mathrm{\text{supp}}\left(\rho \right)\subseteq \left\{\mathit{\theta }\hspace{0.17em}:\hspace{0.17em}{\nabla }_{\mathit{\theta }}\mathrm{\Psi }\left(\mathit{\theta };\rho \right)=\mathbf{0}\hspace{0.17em}\right\}.$$ [20]

Note that global optimizers of $R\left(\rho \right)$, defined by condition (Eq. 17), are fixed points, but the set of fixed points is, in general, larger than the set of optimizers. Our next proposition provides an analogous characterization of the fixed points of diffusion DD (Eq. 12) (see ref. 21 for related results).

Proposition 3:

Assume that conditions A1–A3 hold and that ${\rho }_0$ is absolutely continuous with respect to the Lebesgue measure, with $F_{\beta ,\lambda }\left({\rho }_0\right)<\infty$. If ${\left({\rho }_t\right)}_{t\ge 0}$ is a solution of the diffusion PDE (Eq. 12), then ${\rho }_t$ is absolutely continuous. Further, there is at most one fixed point ${\rho }_{*}={\rho }_{*}^{\beta ,\lambda }$ of Eq. 12 satisfying $F_{\beta ,\lambda }\left({\rho }_{*}\right)<\infty$. This fixed point is absolutely continuous and its density satisfies

$${\rho }_{*}\left(\mathit{\theta }\right)=\frac{1}{Z\left(\beta \right)}\exp \left\{-\beta {\mathrm{\Psi }}_{\lambda }\left(\mathit{\theta };{\rho }_{*}\right)\right\}.$$ [21]

In the next sections, we state our results about convergence of the DD to its fixed point. In the case of noisy SGD [and for the diffusion PDE (12)], a general convergence result can be established (although at the cost of an additional regularization). For noiseless SGD (and the continuity equation; Eq. 12), we do not have such a general result. However, we obtain a stability condition for a fixed point containing one point mass, which is useful to characterize possible limiting points (and is used in treating the examples in the previous section).

Convergence: Noisy SGD.

Remarkably, the diffusion PDE (Eq. 12) generically admits a unique fixed point, which is the global minimum of $F_{\beta ,\lambda }\left(\rho \right)$, and the evolution (Eq. 12) converges to it, if initialized so that $F_{\beta ,\lambda }\left({\rho }_0\right)<\infty$. This statement requires some qualifications. First, we introduce sufficient regularity assumptions to guarantee the existence of sufficiently smooth solutions of Eq. 12:

• i)$V\in C^4\left({\mathbb{R}}^D\right)$, $U\in C^4\left({\mathbb{R}}^D\times {\mathbb{R}}^D\right)$, ${\nabla }_{{\mathit{\theta }}_1}^{k}U\left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right)$ is uniformly bounded for $0\le k\le 4$.

Next notice that the righthand side of the fixed point equation (Eq. 21) is not necessarily normalizable [for instance, it is not when $V\left(\cdot \right)$, $U\left(\cdot ,\cdot \right)$ are bounded]. To ensure the existence of a fixed point, we need $\lambda >0$.

Theorem 4:

Assume that conditions A1–A4 hold, and $1/K_0\le \lambda \le K_0$ for some $K_0>0$. Then $F_{\beta ,\lambda }\left(\rho \right)$ has a unique minimizer, denoted by ${\rho }_{*}^{\beta ,\lambda }$, which satisfies

$$R\left({\rho }_{*}^{\beta ,\lambda }\right)\le \underset{\mathit{\theta }\in {\mathbb{R}}^{N\times D}}{inf}{R}_N\left(\mathit{\theta }\right)+C\hspace{0.17em}D/\beta ,$$ [22]

where $C$ is a constant depending on $K_0,K_1,K_2,K_3$. Further, letting ${\rho }_t$ be a solution of the diffusion PDE (Eq. 12) with initialization satisfying $F_{\beta ,\lambda }\left({\rho }_0\right)<\infty$, we have, as $t\to \infty$,

$${\rho }_t\Rightarrow {\rho }_{*}^{\beta ,\lambda }.$$ [23]

The proof of this theorem is based on the following formula that describes the free-energy decrease along the trajectories of the DD (Eq. 12):

$$\begin{array}{ccc}\hfill & \hfill & \frac{\mathrm{\text{d}}{F}_{\beta ,\lambda }\left({\rho }_t\right)}{\mathrm{\text{d}}t}=\hfill \\ \hfill & \hfill & -2\xi \left(t\right){\int }_{{\mathbb{R}}^D}\Vert {\nabla }_{\mathit{\theta }}\left({\mathrm{\Psi }}_{\lambda }\left(\mathit{\theta };{\rho }_t\right)+1/\beta \cdot \log {\rho }_t\left(\mathit{\theta }\right)\right){\Vert }_2^2{\rho }_t\left(\mathit{\theta }\right)\mathrm{\text{d}}\mathit{\theta }.\hfill \end{array}$$ [24]

(A key technical hurdle is of course proving that this expression makes sense, which we do by showing the existence of strong solutions.) It follows that the righthand side must vanish as $t\to \infty$, from which we prove that (eventually taking subsequences) ${\rho }_t\Rightarrow {\rho }_{*}$ where ${\rho }_{*}$ must satisfy $\beta {\mathrm{\Psi }}_{\lambda }\left(\mathit{\theta };{\rho }_{*}\right)+\log {\rho }_{*}\left(\mathit{\theta }\right)=\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}$. This in turns mean ${\rho }_{*}$ is a solution of the fixed point condition 21 and is in fact a global minimum of $F_{\beta ,\lambda }$ by convexity.

This result can be used in conjunction with Theorem 3, to analyze the regularized noisy SGD algorithm (Eq. 11).

Theorem 5:

Assume that conditions A1–A4 hold. Let ${\rho }_0\in \mathcal{P}\left({\mathbb{R}}^D\right)$ be absolutely continuous with $F_{\beta ,\lambda }\left({\rho }_0\right)<\infty$ and $K_0$ sub-Gaussian. Consider regularized noisy SGD (cf. Eq. 11) at inverse temperature $\beta <\infty$, regularization $1/K_0\le \lambda \le K_0$ with initialization ${\left({\mathit{\theta }}_i^0\right)}_{i\le N}{\sim }_{iid}{\rho }_0$. Then, for any $\eta >0$, there exists $K=K\left(\eta ,\left\{K_i\right\}\right)$, and setting $\beta \ge KD$, there exists $T=T\left(\eta ,V,U,\left\{K_i\right\},D,\beta \right)<\infty$ and $C_0=C_0\left(\eta ,\left\{K_i\right\},\delta \right)$ independent of the dimension $D$ and temperature $\beta$) such that the following happens for $N,\left(1/\epsilon \right)\ge C_0{e}^{C_{0}T}D$, $\epsilon \ge 1/N^{10}$: For any $k\in \left[T/\epsilon ,10T/\epsilon \right]$, we have, with probability $1-\delta$,

$$R_N\left({\mathit{\theta }}^k\right)\le \underset{\rho \in \mathcal{P}\left({\mathbb{R}}^D\right)}{inf}{R}_{\lambda }\left(\rho \right)+\eta .$$ [25]

Let us emphasize that the convergence time $T$ in the last theorem can depend on the dimension $D$ and on the data distribution $\mathbb{P}$ but is independent of the number of hidden units $N$. As illustrated by the examples in the previous section, understanding the dependence of $T$ on $D$ requires further analysis, but examining the proof of this theorem suggests $T=e^{O\left(D\right)}$ quite generally [examples in which $T=O\left(1\right)$ or $T=e^{\mathrm{\Theta }\left(D\right)}$ can be constructed]. We expect that our techniques could be pushed to investigate the dependence of $T$ on $\eta$ (see SI Appendix, Discussion). In highly structured cases, the dimension $D$ can be of constant order and be much smaller than $d$.

Convergence: Noiseless SGD.

The next theorems provide necessary and sufficient conditions for distributions containing a single point mass to be a stable fixed point of the evolution. This result is useful to characterize the large time asymptotics of the dynamics (Eq. 7). Here, we write ${\nabla }_{1}U\left({\mathit{\theta }}_1,{\mathit{\theta }}_2\right)$ for the gradient of $U$ with respect to its first argument and ${\nabla }_{\mathrm{1,1}}^{2}U$ for the corresponding Hessian. Further, for a probability distribution ${\rho }_{*}$, we define

$${\mathbf{H}}_0\left({\rho }_{*}\right)={\nabla }^{2}V\left({\mathit{\theta }}_{*}\right)+\int {\nabla }_{\mathrm{1,1}}^{2}U\left({\mathit{\theta }}_{*},\mathit{\theta }\right)\hspace{0.17em}{\rho }_{*}\left(\mathrm{\text{d}}\mathit{\theta }\right).$$ [26]

Note that ${\mathbf{H}}_0\left({\rho }_{*}\right)$ is nothing but the Hessian of $\mathit{\theta }\mapsto \mathrm{\Psi }\left(\mathit{\theta };{\rho }_{*}\right)$ at ${\mathit{\theta }}_{*}$.

Theorem 6:

Assume $V,U$ to be twice differentiable with bounded gradient and bounded continuous Hessian. Let ${\mathit{\theta }}_{*}\in {\mathbb{R}}^D$ be given. Then ${\rho }_{*}={\delta }_{{\mathit{\theta }}_{*}}$ is a fixed point of the evolution (Eq. 7) if and only if $\nabla V\left({\mathit{\theta }}_{*}\right)+{\nabla }_{1}U\left({\mathit{\theta }}_{*},{\mathit{\theta }}_{*}\right)=\mathbf{0}$.

Define ${\mathbf{H}}_0\left({\delta }_{{\mathit{\theta }}_{*}}\right)\in {\mathbb{R}}^{D\times D}$ as per Eq. 26. If ${\lambda }_{min}\left({\mathbf{H}}_0\left({\delta }_{{\mathit{\theta }}_{*}}\right)\right)>0$, then there exists $r_0>0$ such that, if $\mathrm{\text{supp}}\left({\rho }_{t_0}\right)\subseteq \mathsf{\text{B}}\left({\mathit{\theta }}_{*};r_0\right)\equiv \left\{\mathit{\theta }:\hspace{0.17em}\Vert \mathit{\theta }-{\mathit{\theta }}_{*}{\Vert }_2\le r_0\right\}$, then ${\rho }_t\Rightarrow {\rho }_{*}$ as $t\to \infty$. In fact, convergence is exponentially fast, namely $\int \Vert \mathit{\theta }-{\mathit{\theta }}_{*}{\Vert }_2^2{\rho }_t\left(\mathrm{\text{d}}\mathit{\theta }\right)\le e^{-\lambda \left(t-t_0\right)}$ for some $\lambda >0$.

Theorem 7:

Under the same assumptions of Theorem 6, let ${\rho }_{*}=p_{*}{\delta }_{{\mathit{\theta }}_{*}}+\left(1-p_{*}\right){\tilde{\rho }}_{*}\in \mathcal{P}\left({\mathbb{R}}^D\right)$ be a fixed point of dynamics (Eq. 7), with $p_{*}\in \left(\mathrm{0,1}\right]$ and $\nabla \mathrm{\Psi }\left({\mathit{\theta }}_{*};{\rho }_{*}\right)=\mathbf{0}$ (which, in particular, is implied by the fixed point condition; Eq. 20). Define the level sets $\mathcal{L}\left(\eta \right)\equiv \left\{\mathit{\theta }:\hspace{0.17em}\mathrm{\Psi }\left(\mathit{\theta };{\rho }_{*}\right)\le \mathrm{\Psi }\left({\mathit{\theta }}_{*};{\rho }_{*}\right)-\eta \right\}$ and make the following assumptions: (B1) The eigenvalues of ${\mathbf{H}}_0={\mathbf{H}}_0\left({\rho }_{*}\right)$ are all different from 0, with ${\lambda }_{min}\left({\mathbf{H}}_0\right)<0$; (B2) ${\tilde{\rho }}_{*}\left(\mathcal{L}\left(\eta \right)\right)\uparrow 1$ as $\eta \downarrow 0$; and (B3) there exists ${\eta }_0>0$ such that the sets $\partial \mathcal{L}\left(\eta \right)$ are compact for all $\eta \in \left(0,{\eta }_0\right)$.

If ${\rho }_0$ has a bounded density with respect to the Lebesgue measure, then it cannot be that ${\rho }_t$ converges weakly to ${\rho }_{*}$ as $t\to \infty$.

Discussion and Future Directions

In this paper, we developed an approach to the analysis of two-layer neural networks. Using a propagation-of-chaos argument, we proved that—if the number of hidden units satisfies $N\gg D$—SGD dynamics is well approximated by the PDE in Eq. 7, while noisy SGD is well approximated by Eq. 12. Both of these asymptotic descriptions correspond to Wasserstein gradient flows for certain energy (or free energy) functionals. While empirical risk minimization is known to be insensitive to overparametrization (22), the present work clarifies that the SGD behavior is also independent of the number of hidden units, as soon as this is large enough.

We illustrated our approach on several concrete examples, by proving convergence of SGD to a near-global optimum. This type of analysis provides a mechanism for avoiding the perils of nonconvexity. We do not prove that the finite-$N$ risk $R_N\left(\mathit{\theta }\right)$ has a unique local minimum or that all local minima are close to each other. Such claims have often been the target of earlier work but might be too strong for the case of neural networks. We prove instead that the PDE (Eq. 7) converges to a near-global optimum, when initialized with a bounded density. This effectively gets rid of some exceptional stationary points of $R_N\left(\mathit{\theta }\right)$ and merges multiple finite $N$ stationary points that result into similar distributions $\rho$.

In the case of noisy SGD (Eq. 11), we prove that it converges generically to a near-global minimum of the regularized risk, in time independent of the number of hidden units.

We emphasize that while we focused here on the case of square loss, our approach should be generalizable to other loss functions as well (cf. SI Appendix).

The present work opens the way to several interesting research directions. We will mention two of them: $\left(i\right)$ The PDE (Eq. 7) corresponds to gradient flow in the Wasserstein metric for the risk $R\left(\rho \right)$ (see ref. 20). Building on this remark, tools from optimal transportation theory can be used to prove convergence. $\left(ii\right)$ Multiple finite-$N$ local minima can correspond to the same minimizer ${\rho }_{*}$ of $R\left(\rho \right)$ in the limit $N\to \infty$. Ideas from glass theory (24) might be useful to investigate this structure.

Let us finally mention that, after a first version of this paper appeared as a preprint, several other groups obtained results that are closely related to Theorem 3 (25–27).