Analysis of a Two-Layer Neural Network via Displacement Convexity

Introduction

One of the most fruitful ideas in this context is to use functions that are linear combinations of simple components:

Special instantiations of this idea include (we provide only pointers to the immense literature on each topic):

Kernel ridge regression and related random feature methods [CST00, RR08];

Despite the impressive practical success of these methods, the risk function $R_{N}({\boldsymbol{w}})$ is highly non-convex and little is known about global convergence of algorithms that try to minimize it (we refer to Section 2 for further discussion of the related literature).

Notable exceptions to the last statement are provided by random features and by boosting algorithms. In random feature methods, the parameters ${\boldsymbol{w}}_{i}$ are not optimized over (they are drawn i.i.d. from some common distribution), and the resulting risk function becomes convex in the weights $(a_{1},\dots,a_{N})$ to be learnt. While this is a fruitful idea, it gives up the degrees of freedom afforded by the ${\boldsymbol{w}}_{i}$’s.

Boosting overcomes non-convexity by fitting the components ${\boldsymbol{w}}_{1}$, …, ${\boldsymbol{w}}_{N}$ one at the time, sequentially. The underlying assumption is that the problem of minimizing $R_{N}({\boldsymbol{w}})$ with respect to one of the hidden units ${\boldsymbol{w}}_{i}$ is tractable. However, this is generally not the case when the parameters ${\boldsymbol{w}}_{i}$ belong to a high-dimensional space.

The risk function (1.2) crystalizes a central conundrum in statistical learning. In a number of applications (especially at low noise), it is rarely the case that low prediction error can be achieved through a function that is linear in the raw covariates, e.g. $\hat{f}(x)=\langle{\boldsymbol{w}},{\boldsymbol{x}}\rangle$. In a classical setting, the statistician would craft nonlinear features out of the covariates on the basis of expert knowledge. For the model of Eq. (1.1), this amounts to constructing vectors ${\boldsymbol{w}}_{1},\dots,{\boldsymbol{w}}_{N}$. Statistical methods would then be confined to the convex task of fitting the coefficients $a_{1},\dots,a_{N}$. This step is well understood from a statistical and computational perspective.

Modern machine learning approaches (boosting, neural networks, etc.) hold the promise of automatizing feature extraction, hence producing superior performances in a wide variety of applications. Unfortunately, we are still far from understanding in which cases optimizing over the ${\boldsymbol{w}}_{i}$’s yields a significant improvement over –say– choosing them randomly. This central challenge intertwines statistical and computational aspects. It is not hard to see that varying the weights ${\boldsymbol{w}}_{i}$’s produces a significantly larger function class [Bac17]. The relevant question is what part of this class can be accessed using gradient descent or other practical algorithms.

The main objective of this paper is to introduce a nonparametric regression model in which these questions can be addressed rigorously. The model is interesting for at least two reasons: $(i)$ From a theoretical point of view, global convergence can be proved in the limit of a large neurons. The proof relies on a mathematical mechanism that has not been explored in the statistics or machine learning literature before. $(ii)$ From a practical point of view, the model is nontrivial enough to illustrate the potential advantage of fitting the features ${\boldsymbol{w}}_{i}$ (we demonstrate this numerically in Section 4.)

The model (1.4) is general enough to include a broad class of radial-basis function (RBF) networks which are known to be universal function approximators [PS91]. To the best of our knowledge, there is no result on the global convergence of stochastic gradient descent for learning RBF networks, and this paper establishes the first result of this type.

It is important to emphasize a few differences with respect to standard RBF networks. First of all, we do not require the kernel $K({\boldsymbol{x}})$ to be radial, i.e. to depend uniquely on the norm $|{\boldsymbol{x}}|$. Second, we require $K$ to have compact support. This is mainly a technical requirement that simplifies some arguments: we expect our results to be generalizable to kernels that decay rapidly enough. Finally, and most crucially, the form (1.4) does not include non-uniform weights for the $N$ components. A more standard formulation would posit $\hat{f}({\boldsymbol{x}};{\boldsymbol{w}})=\sum_{i=1}^{N}a_{i}K^{\delta}({\boldsymbol{x}}-{\boldsymbol{w}}_{i})$ and learn the weights $a_{i}$ from data, see Eq. (1.1). We deliberately set the weights to a fixed value because the risk function is convex in ${\boldsymbol{a}}=(a_{i})_{i\leq N}$, and hence fitting ${\boldsymbol{a}}$’s to global optimality is ‘easy.’ Indeed, universal approximation could be achieved by keeping the centers ${\boldsymbol{w}}_{i}$ fixed (and sufficiently dense in $\Omega$) and only adjusting ${\boldsymbol{a}}$. As discussed above, our focus is on the role of the ${\boldsymbol{w}}_{i}$’s.

Our main result is a proof that, for sufficiently large $N$ and small $\delta$, gradient descent algorithms converge to weights ${\boldsymbol{w}}$ with nearly optimum prediction error, provided $f$ is strongly concave. Let us emphasize that the resulting population risk $R_{N}({\boldsymbol{w}})$ is non-convex regardless of the concavity properties of $f$. Our proof unveils a novel mechanism by which global convergence takes place. Convergence results for non-convex empirical risk minimization are generally proved by carefully ruling out local minima in the cost function (see Section 2 for pointers to this literature). Instead we prove that, as $N\to\infty$, $\delta\to 0$, the gradient descent dynamics converges to a gradient flow in Wasserstein space, and that the corresponding cost function is ‘displacement convex.’ Breakthrough results in optimal transport theory guarantee dimension-free convergence rates for this limiting dynamics [CJM+01, CMV03, CMV06]. In particular, we expect the cost function $R_{N}({\boldsymbol{w}})$ to have many local minima, which are however completely neglected by the gradient descent dynamics.

More specifically, our first step is to show that – for large $N$ – the evolution of the weights ${\boldsymbol{w}}_{1},\dots,{\boldsymbol{w}}_{N}$ under gradient descent can be replaced by the evolution of a probability distributionThroughout,$\mathscrsfs{P}_{2}({\cal X})$ denotes the space of probability distributions on ${\cal X}$, endowed with Wasserstein metric $W_{2}$. $\rho^{\delta}\in\mathscrsfs{P}_{2}(\Omega)$, which approximates their empirical distribution. Namely, if $({\boldsymbol{w}}^{k}_{1},\dots,{\boldsymbol{w}}^{k}_{N})$ denote the weights after $k$ iterations with step size ${\varepsilon}$, and $\hat{\rho}^{(N)}_{k}=\sum_{i=1}^{N}\delta_{{\boldsymbol{w}}_{i}^{k}}/N$ is their empirical distribution, then we have

where the limit holds in the sense of weak convergence or in $W_{1}$ distance (the two are equivalent since $\Omega$ is compact). The limit evolution $(\rho^{\delta}_{t})_{t\geq 0}$ satisfies a partial differential equation (PDE) that can also be described as the Wasserstein $W_{2}$ gradient flow (i.e. gradient flow in $\mathscrsfs{P}_{2}(\Omega)$), for the following effective risk

where $\nu_{0}=1/|\Omega|$ and $|\Omega|$ denotes the volume of the set $\Omega$. Here $\ast$ denotes the usual convolution. Let us emphasize that the convergence to Wasserstein gradient flow holds regardless of the concavity of $f$.

The use of $W_{2}$ gradient flows to analyze two-layer neural networks was recently developed in several papers [MMN18, RVE18, CB18, SS18]. However, we cannot rely on earlier results because of the specific boundary conditions in our problem. We constrain the ${\boldsymbol{w}}_{i}\in\Omega^{\delta}$ by running projected stochastic gradient descent (SGD): at each step ${\boldsymbol{w}}_{i}$ moves in the direction of a stochastic gradient of $R_{N}({\boldsymbol{w}})$ and then projected back to $\Omega^{\delta}$. This results in a PDE with Neumann boundary condition on $\Omega^{\delta}$, which is not covered by previous theory. We establish a quantitative version of the limit (1.5) via propagation-of-chaos techniques.

Even if the cost (1.6) is quadratic and convex in $\rho$, its $W_{2}$ gradient flow can have multiple fixed points, and hence global convergence cannot be guaranteed. Global convergence results were proven in [MMN18] and in [CB18] by showing that, for all $t\geq 0$ $\rho^{\delta}_{t}$ has a density that is either smooth, or strictly positive everywhere. However, these convergence results are non-quantitative, and do not provide convergence ratesAn argument indicating convergence in a time polynomial in $d$ was put forward in [WLLM18], but for a different type of continuous flow..

Indeed, the mathematical property that controls global convergence of $W_{2}$ gradient flow is not ordinary convexity but displacement convexity. Roughly speaking, displacement convexity is convexity along geodesics of the $W_{2}$ metric, see Section 3.5. The risk function (1.6) is not displacement convex. Indeed, its quadratic term reads $\nu_{0}\int K_{\delta}\ast K_{\delta}({\boldsymbol{x}}-{\boldsymbol{x}}^{\prime})\rho({\boldsymbol{x}})\rho({\boldsymbol{x}}^{\prime}){\rm d}{\boldsymbol{x}}{\rm d}{\boldsymbol{x}}^{\prime}$ which is not displacement convex unless $K_{\delta}\ast K_{\delta}$ is convex (see Lemma H.1), which cannot be in our setting. However, for small $\delta$, we can formally approximate $K^{\delta}\ast\rho\approx\rho$, and hence hope to replace the risk function (1.6) with a simpler one

Most of our technical work is devoted to making rigorous this $\delta\to 0$ approximation. Namely, we prove that, as $\delta\to 0$, $\rho^{\delta}_{t}\Rightarrow\rho_{t}$ where $\rho_{t}$ follows the $W_{2}$ gradient flow for the risk $R(\rho)$.

Remarkably, the risk function $R(\rho)$ is strongly displacement convex (provided $f$ is strongly concave). A long line of work in PDE and optimal transport theory establishes dimension-free convergence rates for its $W_{2}$ gradient flow [CJM+01, CMV03, CMV06]. Namely, if $f$ is $\alpha$-strongly concave, then $R(\rho_{t})\leq R(\rho_{0})\,e^{-2\alpha t}$. By using the approximation results outlined above, we obtain global convergence for SGD. With high probability,

where the error term ${\sf err}$ vanishes as $N\to\infty$, ${\varepsilon},\delta\to 0$ in a suitable order.

This result implies that SGD converges exponentially fast to a near-global optimum with a rate that is controlled by the convexity parameter $\alpha$.

Our bounds are not sharp enough to provide quantitative control on the error term ${\sf err}(N,d,{\varepsilon},\delta)$, especially in high dimension. Nevertheless, the convergence rate predicted by our asymptotic theory is in excellent agreement with numerical simulations, cf. Section 4. Explaining this surprising quantitative agreement is an outstanding challenge.

Related literature

The present work ties in several lines of research, some of which were already mentioned in the introduction. A substantial amount of work has been devoted to analyzing two-layer neural networks and developing algorithms with convergence guarantees, see e.g. [ZSJ+17, Tia17, BJW18]. However these approaches are typically based on tensor factorization or similar initialization steps that are not used in practice, and do not scale well (although polynomially) in high dimension.

The landscape of empirical risk minimization was also studied in a number of papers, see e.g. [LY17, SJL18]. However, global convergence was only proved in the extremely overparametrized regime in which the neural network essentially behaves as kernel ridge regression [DZPS18].

Classical theory of neural networks was largely devoted to the two-layer case [AB09], although the focus was on representation and approximation questions [Cyb89, Bar93], as well as on generalization error. It was already clear in that context that a two-layer network is conveniently characterized by the empirical distribution of the hidden neurons, and that it is useful to relax this from a distribution with $N$ atoms, to a general probability measure. This representation plays an important role, for instance, in [Bar98], and was exploited again under the label of ‘convex neural networks’ in [BRV+06].

Over the last year, several groups independently revisited this connection, with the objective of understanding the landscape structure of two-layer networks, and the dynamics of gradient descent methods [NS17, MMN18, RVE18, SS18, CB18, MMM19]. In particular, it was proven in [MMN18] that, under certain smoothness condition on the underlying data distribution, the gradient descent evolution is well approximated by a Wasserstein gradient flow, provided that the number of neurons exceeds the data dimensions. As mentioned above, the algorithm treated here differs from the ones analyzed in earlier work, because the weights ${\boldsymbol{w}}_{i}$ are constrained to lie in the convex set $\Omega^{\delta}$. We enforce this constraint by using projected SGD, i.e. projecting at each step the weights onto the set $\Omega^{\delta}$. We generalize the analysis of [MMN18], obtaining convergence to a PDE with Neumann (reflecting) boundary conditions. As in [MMN18], we build on ideas that were first developed in the context of interacting particle systems [Dob79, Szn91].

The Wasserstein gradient flow approach was used in [MMN18, CB18] to establish global convergence results. However, these results fall short of our objectives for several reasons:

The global convergence result of [CB18] rely on certain homogeneity properties of the neurons that are lacking here. We could obtain homogeneity by adding coefficients to Eq. (1.4), i.e. considering $\hat{f}({\boldsymbol{x}};{\boldsymbol{w}})=\sum_{i=1}^{N}a_{i}K^{\delta}({\boldsymbol{x}}-{\boldsymbol{w}}_{i})$ and minimizing the risk with respect to the coefficients $a_{i}$. As mentioned above, we refrain from introducing coefficients not to oversimplify the problem: when $N\to\infty$, it is sufficient to fit the coefficients $a_{i}$ to achieve vanishing risk. Fitting the $a_{i}$’s is a least squares problem.

In contrast, our results are a first step towards dimension-independent convergence rate, in a more restricted setting than [MMN18, CB18, MMM19].

In summary, our results do not subsume earlier work, that assumes a more general setting, but rather establish stronger results in narrower context. Indeed, we believe that specific structural conditions must be imposed on the data distribution and activation function for the Wasserstein gradient flow approach to yield quantitative convergence rates. This paper presents one specific set of assumptions. Although our results are not strong enough to establish non-asymptotic convergence rates, they point clearly in that direction.

Model and assumptions

Throughout the paper, we use $C$ to denote finite constants, which can vary from point to point. When these constants can depend on some of the problem parameters, e.g. $a,b,c$, we will write $C(a,b,c)$. When they are absolute numerical constants, we will emphasize this by writing $C_{*}$.

2 Data

Our formal assumptions on the set $\Omega$ and the function $f$ are as follows:

$\Omega\supseteq{\sf B}({\boldsymbol{0}};r)$, with $r>0$, is a compact convex set with $\mathscrsfs{C}^{2}$ boundary.

where $\nabla^{2}f$ denotes the Hessian of $f$.

$f\in\mathscrsfs{C}^{\infty}(\Omega)$, with $\|f\|_{\mathscrsfs{L}^{\infty}(\Omega)},\|\nabla f\|_{\mathscrsfs{L}^{\infty}(\Omega)}\leq C_{*}$ for an absolute constant $C_{*}$.

Without loss of generality, we can also assume that $\int_{\Omega}f({\boldsymbol{x}})\,{\rm d}{\boldsymbol{x}}=1$. As a running example, we will use $\Omega={\sf B}({\boldsymbol{0}};r)$, where we remind $r$ is defined in Assumption (A1).

The assumption ${\boldsymbol{x}}_{j}\sim{\sf Unif}(\Omega)$ is quite strong but simplifies our analysis. We believe our approach can be generalized to a broader family of probability distribution for the covariates ${\boldsymbol{x}}_{j}$, but defer these generalizations to future work.

3 Neural network and SGD

For $\delta>0$, let $K^{\delta}({\boldsymbol{x}})=\delta^{-d}K({\boldsymbol{x}}/\delta)$. We try to fit the function (1.4) with parameters ${\boldsymbol{w}}=({\boldsymbol{w}}_{1},\dots,{\boldsymbol{w}}_{N})$. These parameters are constrained to ${\boldsymbol{w}}_{i}\in\Omega^{\delta}$ which is a suitable scaling of $\Omega$, as defined in the following. Given $\delta<r/c_{0}$, with $r$ defined in (A1), define

We use stochastic gradient descent to minimize the population risk (1.2). At each step, we use a new data point $(y_{k},{\boldsymbol{x}}_{k})$, thus the sample size is equal to the number of iterations of the algorithm. Assuming for simplicity constant step size ${\varepsilon}>0$, we update the parameters by

Here ${\boldsymbol{g}}_{i}^{k+1}\sim{\sf N}(0,{\boldsymbol{I}}_{d})$ is Gaussian noise which we take to be i.i.d. across time and neuron indices, $k$ and $i$, and ${\sf P}$ is the orthogonal projector onto $\Omega^{\delta}$:

The noise term $\sqrt{2{\varepsilon}\tau}\,{\boldsymbol{g}}_{i}^{k+1}$ is added mainly for technical reasons. Namely, it allows us to control the smoothness of the solutions of the resulting PDE. In simulations we do not find it useful, and we believe that a more careful analysis would be able to establish smoothness without the noise term.

We initialize SGD with $({\boldsymbol{w}}_{i}^{0})_{i\leq N}\sim_{\rm i.i.d.}\rho^{\delta}_{\mbox{\tiny\rm init}}\in\mathscrsfs{P}_{2}(\Omega^{\delta})$, where $\rho^{\delta}_{\mbox{\tiny\rm init}}$ is a scaling of a fixed distribution $\rho_{\mbox{\tiny\rm init}}\in\mathscrsfs{P}_{2}(\Omega)$, i.e. $\rho^{\delta}_{\mbox{\tiny\rm init}}(S)=\rho_{\mbox{\tiny\rm init}}(S/\lambda_{\delta})$. We assume that the initialization is smooth:

$\rho_{\mbox{\tiny\rm init}}\in\mathscrsfs{C}^{\infty}(\Omega^{\delta})$.

4 PDE Model, δ>0𝛿0\delta>0

In particular $\lim_{\delta\to 0}\inf_{\rho\in\mathscrsfs{P}_{2}(\Omega)}R^{\delta}(\rho)=0$.

Our first main result is that the dynamics of SGD is well approximated by the following PDE (see Section 5.1 for a formal statement):

where ${\boldsymbol{n}}({\boldsymbol{x}})$ denotes the inward normal vector to $\partial\Omega^{\delta}$ at ${\boldsymbol{x}}$.

A rigorous definition of solutions of this PDE, along with some of their properties, is given in Appendix B. In Appendix C, we discuss the connection between the PDE (3.9) and the so-called “nonlinear dynamics”, i.e. a stochastic differential equation that captures the trajectories of the weights ${\boldsymbol{w}}_{i}^{k}$. Using this connection, we prove existence and uniqueness of weak solutions of Eq. (3.9). In the proofs, we will often assume $\nu_{0}=1$, which amounts to a rescaling of time $t$.

For $\tau=0$, the evolution defined by Eq. (3.9) corresponds to the gradient flow in Wasserstein metric for the risk function $R^{\delta}(\rho)$. For $\tau>0$, it is the gradient flow for the free energy functional $F^{\delta}(\rho)$ defined below

5 Limit PDE, δ=0𝛿0\delta=0

The corresponding Wasserstein gradient flow is also known as viscous porous medium equation [Váz07] and it is given by

In Appendix A, we give the definition of a weak solution for the PDE (3.12) with initial and boundary conditions (3.13). We also prove that the weak solution of the PDE (3.12) is unique, under a mild integrability condition. Again, in proofs we will assume without loss of generality $\nu_{0}=1$.

As in the $\delta>0$ case, the evolution defined by Eq. (3.12) is the gradient flow for the free energy $F(\rho)=(1/2)R(\rho)-\tau S(\rho)$. Our analysis uses a key property of the risk function $R(\rho)=\nu_{0}\|f-\rho\|_{\mathscrsfs{L}^{2}(\Omega)}^{2}$ (and the free energy): displacement convexity [McC97]. For the reader’s convenience, we recall its definition here, referring to [AGS08, Vil08, San15] for further background. Given two probability measures $\rho_{0},\rho_{1}\in\mathscrsfs{P}_{2}(\Omega)$, their $W_{2}$ distance is defined by

where the infimum is taken over the set $\Gamma(\rho_{0},\rho_{1})$ of couplings of $\rho_{0}$, $\rho_{1}$ (i.e. probability measures on $\Omega\times\Omega$ whose first marginal coincides with $\rho_{0}$, and second with $\rho_{1}$). The infimum is achieved by weak compactness of $\mathscrsfs{P}_{2}(\Omega)$.

The metric space $(\mathscrsfs{P}_{2}(\Omega),W_{2})$ is a ‘length space,’ and in particular it is possible to construct geodesics, i.e. paths of minimum length connecting any two probability measures $\rho_{0},\rho_{1}$. Geodesics have a simple description. Let $\gamma_{*}$ be the coupling achieving the infimum in the definition of $W_{2}(\rho_{0},\rho_{1})$. Letting $(\boldsymbol{X}_{0},\boldsymbol{X}_{1})\sim\gamma_{*}$, we define $\rho_{t}$ to be the distribution of $\boldsymbol{X}_{t}=(1-t)\boldsymbol{X}_{0}+t\boldsymbol{X}_{1}$. The curve $t\mapsto\rho_{t}$, indexed by $t\in$ turns out to be the geodesic between $\rho_{0}$ and $\rho_{1}$ in $(\mathscrsfs{P}_{2}(\Omega),W_{2})$.

A useful observation is that displacement convexity implies that all local minima of ${\mathcal{F}}$ are global minimizer. Indeed, by (3.15) it is straightforward to see that ${\mathcal{F}}$ has at most one global minimizer $\rho^{*}$. Also, for every other point $\rho$, the geodesic between $\rho$ and $\rho_{*}$ is a strictly decreasing path for the function ${\mathcal{F}}$. Now, suppose that $\bar{\rho}\neq\rho_{*}$ is a local minimum. Then, there exists a neighborhood $U$ around $\bar{\rho}$ such that, for any $\rho\in U$, ${\mathcal{F}}(\rho)\geq{\mathcal{F}}(\bar{\rho})$. However, the strictly decreasing path between $\bar{\rho}$ and $\rho_{*}$ passes through the neighborhood $U$, which leads to a contradiction and so $\rho=\rho_{*}$

It follows from [McC97] that the risk function $R(\rho)$ and the free energy $F(\rho)$ are strongly displacement convex.

The concavity assumption on the regression function $f$ (Assumption (A2)) defines a nonparametric class under which global convergence can be established, with convergence rates uniquely determined by the curvature $\alpha$ (in the limit $N\to\infty$, $\delta\to 0$). Nonparametric estimation of concave functions has attracted considerable attention over recent years, see e.g. [HD13, CS16], and is –by itself– an interesting domain of applicability.

However, our projected SGD algorithm is potentially applicable to any data set, and will return a meaningful estimate $\hat{f}$ regardless whether $f$ is concave or not. Indeed, in the next section we present numerical simulations indicating convergence to a near-global optimum even for non-concave functions $f$.

From mathematical point of view, Assumption (A2) is only used to show the convergence of the solution of the viscous porous medium equation (limit PDE, $\delta=0$) to the unique global minimizer of the free energy $F(\rho)=(1/2)R(\rho)-\tau S(\rho)$, as formally stated in Theorem F.8. Concavity is not needed for the other results in the paper, namely approximating the SGD trajectory with the solution of the PDE ($\delta>0$), see Theorem 5.1, and the convergence of the solution of the PDE ($\delta>0$) to the solution of the viscous porous medium equation, see Theorem 5.2. It is therefore foreseeable a more general analysis that relaxes the concavity assumption.

Numerical illustrations

In this section we provide some simple numerical illustrations of our setting, and compare numerical results with the predictions of the Wasserstein gradient flow theory.

We set $\Omega=$ and $f(x)=(1-e^{x-1})/(1-e^{-2})$ (we choose the normalization so that $\int_{-1}^{1}f(x){\rm d}x=1$). Note that $f$ is uniformly concave in $$. We set the kernel$K$ as follows:

where $C_{d}$ is a normalization constant ensuring that $\int_{-1}^{1}K(x)\,{\rm d}x=1$. The initialization $\rho_{\mbox{\tiny\rm init}}$ is a truncated Gaussian: $\rho_{\mbox{\tiny\rm init}}(x)=c\cdot\exp(-x^{2}/(2\sigma^{2}))\,{\boldsymbol{1}}_{}(x)$, with $\sigma=1/3$.

We find empirically that standard stochastic gradient descent (SGD) without the projection ${\sf P}$ onto $\Omega^{\delta}$ works well in this example, and consider this algorithm for simplicity in our first illustrations. We pick $N=200$, $\tau=0$ (noiseless SGD), and constant step size $\varepsilon=10^{-6}$. In Figure 1, left column, we plot the true function $f(\,\cdot\,)$ together with the neural network estimate $\hat{f}(\,\cdot\,;{\boldsymbol{w}}^{k})$ at several points in time $t$ (time is related to the number of iterations $k$ via $t=k{\varepsilon}$). Different plots correspond to different values of $\delta$ with $\delta\in\{1/5,1/10,1/20\}$. We observe that the network estimates $\hat{f}(\,\cdot\,;{\boldsymbol{w}}^{k})$ seem to converge to a limit curve which is an approximation of the true function $f$. As expected, the quality of the approximation improves as $\delta$ gets smaller.

In the right column, we report the evolution of the population risk (1.2) normalized by $\|f\|^{2}_{\mathscrsfs{L}^{2}(\Omega)}$. For comparison, we plot the evolution of the risk (1.7) as predicted by the limit PDE (3.12) with $\tau=0$. We solve the PDE (3.12) numerically using a finite difference scheme that enforces the conservation law $\int\rho(x,t){\rm d}x=1$, see, e.g., [Tho13]. In the finite difference scheme, we choose time step and spatial step $\Delta t=10^{-5}$ and $\Delta x=10^{-2}$, respectively. The curve obtained by this numerical solution appears to capture well the evolution of SGD towards optimality. The main difference is that, while the PDE (3.12) corresponds to $\delta=0$, and hence evolves towards a global optimum at zero risk, SGD converges to a non-zero risk value, which can be interpreted as the approximation error, decreasing with $\delta$.

In Figure 2, we illustrate the numerical solution of the PDE (3.12) by plotting (i) the regression function $f$ together with the PDE solution $\rho_{t}$ (which coincides with the prediction $\hat{f}$ at $\delta=0$) at several times $t$, and (ii) the PDE prediction for the risk $R(\rho_{t})$ (1.7) normalized with respect to $\|f\|^{2}_{\mathscrsfs{L}^{2}(\Omega)}$ (this plot aggregates data from Figs. 1.(b), (d), (f)). We also compare the risk (1.7) to the population risk $R_{N}({\boldsymbol{w}}^{k})$ achieved by SGD for different values of $\delta$. Note that, as $\delta$ becomes smaller, the risk converges to the predicted curve. The risk of the limit PDE (3.12) converges to exponentially fast in $t$, as predicted by the strong displacement convexity of $R(\rho)$.

In Figure 3, we consider the SGD algorithm with projection ${\sf P}$, see (3.5). We pick $N=200$, $\tau=0$, ${\varepsilon}=10^{-6}$ and $\delta=1/20$. On the left, we illustrate the evolution of the value of $40$ weights chosen at random; and on the right, we plot the histogram of their empirical distribution at $t=5$. Note that this histogram matches well the regression function $f$ plotted in black.

2 A two-dimensional concave example

Next, we consider a two-dimensional example. We set $\Omega=^{2}$ and

with ${\boldsymbol{q}}_{1}=(2.5127,-2.4490)$, ${\boldsymbol{q}}_{2}=(0.0596,1.9908)$ and where $c_{1}$ and $c_{2}$ are chosen so that $f$ is non-negative and $\int_{\Omega}f({\boldsymbol{x}})\,{\rm d}{\boldsymbol{x}}=1$. The kernel $K$ is given by $K({\boldsymbol{x}})=C_{d}\kappa(|{\boldsymbol{x}}|)$, where $\kappa$ is defined in (4.1) and $C_{d}$ is a normalization constant ensuring that $\int_{{\sf B}({\boldsymbol{0}};1)}K({\boldsymbol{x}})\,{\rm d}{\boldsymbol{x}}=1$. Again, the initialization $\rho_{\mbox{\tiny\rm init}}$ is a truncated Gaussian: $\rho_{\mbox{\tiny\rm init}}({\boldsymbol{x}})=c\cdot\exp(-|{\boldsymbol{x}}|^{2}/(2\sigma^{2}))\,{\boldsymbol{1}}_{^{2}}({\boldsymbol{x}})$, with $\sigma=1/3$. We compare the normalized risk of SGD with no projection ${\sf P}$ ($N=2000$, $\tau=0$ and ${\varepsilon}=10^{-6}$) for $\delta\in\{1/3,1/5,1/10\}$ with that of the limit PDE (3.12). Figure 4 shows that, already at $\delta=1/10$, the risk of SGD converges to the predicted curve and the risk of the limit PDE (3.12) tends to exponentially fast in $t$.

3 Comparing feature learning to random features

As discussed in the introduction, it is useful to consider the more general model

with parameters ${\boldsymbol{a}}=(a_{1},\ldots,a_{N})$ as well as ${\boldsymbol{w}}=({\boldsymbol{w}}_{1},\dots,{\boldsymbol{w}}_{N})$. This setting allows to compare two different approaches:

Random feature regression: the weights ${\boldsymbol{w}}$ are chosen independently of the labels $y_{i}$ (we allow for dependence on the covariates ${\boldsymbol{x}}_{i}$).

Feature learning: the weights ${\boldsymbol{w}}$ depend on the data $(y_{i},{\boldsymbol{x}}_{i})$.

where $\lambda$ is chosen via cross-validation on a hold-out set, comprising $10\%$ of the samples.

In Figure 5, we compare the performance of three different ways to construct the weights ${\boldsymbol{w}}$: ‘random ${\boldsymbol{w}}$,’ we choose the weights ${\boldsymbol{w}}_{i}$ independently and uniformly at random in $\Omega$ (blue triangles pointing down); ‘${\boldsymbol{w}}=$ data points,’ we choose the weights ${\boldsymbol{w}}_{i}$ uniformly at random among the data points (green circles); ‘optimized ${\boldsymbol{w}}$,’ we use the output of the projected SGD algorithm of the previous sections (red triangles pointing up). The first two can be regarded as ‘random features’ approaches, while the latter is a ‘feature learning’ method.

For the optimized ${\boldsymbol{w}}$, we use exactly the same algorithm in as in (3.5) (without coefficients ${\boldsymbol{a}}$ in the SGD update), with the only difference that each SGD step is carried out with respect to an independent sample from the empirical data, with replacement. SGD is stopped after $k_{\max}$ iteration, and the coefficient $\hat{{\boldsymbol{a}}}$ are computed according to (4.3). Notice that this procedure is probably suboptimal, and it would be better to optimize ${\boldsymbol{a}}$ and ${\boldsymbol{w}}$ jointly: we choose this simpler two-stage procedure to have a more direct application of the algorithm analyzed in the paper, and a comparison with the random feature methods. We set $\tau=0$ (noiseless SGD), and constant step size $\varepsilon=5\cdot 10^{-4}$. The number of iterations $k_{\max}\in\{5\cdot 10^{3},15\cdot 10^{3},5\cdot 10^{4},15\cdot 10^{4},5\cdot 10^{5},15\cdot 10^{5}\}$ is chosen via cross-validation, by using the same hold-out set employed to optimize $\lambda$.

We set $\Omega=^{4}$ and define $y_{j}=f({\boldsymbol{x}}_{j})$, where $f({\boldsymbol{x}})$ takes the form (4.2) with ${\boldsymbol{q}}_{1}=(-0.3832,0.3074,-0.3198,0.4792)$ and ${\boldsymbol{q}}_{2}=(0.3502,-0.1471,$ $0.1685,0.0546)$. Again, $c_{1}$ and $c_{2}$ are chosen so that $f$ is non-negative and $\int_{\Omega}f({\boldsymbol{x}})\,{\rm d}{\boldsymbol{x}}=1$; the kernel $K$ is given by $K({\boldsymbol{x}})=C_{d}\kappa(|{\boldsymbol{x}}|)$, where $\kappa$ is defined in Eq. (4.1) and $C_{d}$ ensures that $\int_{{\sf B}({\boldsymbol{0}};1)}K({\boldsymbol{x}})\,{\rm d}{\boldsymbol{x}}=1$.

After estimating ${\boldsymbol{w}}_{i}$ and $a_{i}$ by either methods, we generate a test set of $10,000$ samples and use it to estimate the generalization error. We perform $20$ independent trials of the experiment, and we plot the average risk normalized by $\|f\|^{2}_{\mathscrsfs{L}^{2}(\Omega)}$ together with the error bar at 1 standard deviation. In Figure 5-(a), we fix the number of neurons $N=200$ and we plot the normalized risk as a function of the number of data points $n$. In Figure 5-(b), we fix the number of samples $n$ to $2000$ and we plot the normalized risk as a function of the number of neurons $N$. The data set used for cross-validation has size $\max(n/10,40)$. Note that feature learning leads to improved performance in both settings. The improvement becomes more pronounced with the sample size $n$, presumably because a better set of weights ${\boldsymbol{w}}_{i}$ can be learnt. On the other hand, when the number of neurons $N$ becomes very large, random ${\boldsymbol{w}}_{i}$’s are already covering $\Omega$ densely enough, and there is no significant advantage in feature learning.

4 A non-concave one-dimensional example

We set $\Omega=$ and $f(x)=(x+\sin(5x-\pi/2)-c_{1})/c_{2}$, where $c_{1}$ and $c_{2}$ are chosen so that $f$ is non-negative and $\int_{\Omega}f(x)\,{\rm d}x=1$. Note that the target function $f$ is bimodal, thus it is not concave. We perform the same numerical experiment described in Section 4.1. In Figure 6, left column, we plot the true function $f(\,\cdot\,)$ together with the neural network estimate $\hat{f}(\,\cdot\,;{\boldsymbol{w}}^{k})$ at several points in time $t$, where different plots correspond to different values of $\delta\in\{1/5,1/10,1/20\}$. In the right column, we report the evolution of the population risk (1.2) normalized by $\|f\|^{2}_{\mathscrsfs{L}^{2}(\Omega)}$. In Figure 7, we plot (i) the regression function $f$ together with the PDE solution $\rho_{t}$ at several times $t$, and (ii) the PDE prediction for the risk $R(\rho_{t})$ (1.7) (normalized with respect to $\|f\|^{2}_{\mathscrsfs{L}^{2}(\Omega)}$) compared with the population risk $R_{N}({\boldsymbol{w}}^{k})$ achieved by SGD for different values of $\delta$. Even if the target function is not concave, the results are similar to those presented in the concave case: (i) the network estimates $\hat{f}(\,\cdot\,;{\boldsymbol{w}}^{k})$ seem to converge to a limit curve which is an approximation of the true function $f$, (ii) the quality of the approximation improves as $\delta$ gets smaller, and (iii) the risk of the limit PDE (3.12) converges to exponentially fast in $t$.

5 Failure for small N𝑁N

We repeat the same experiment described in Section 4.1 for a smaller number of neurons $N=20$. As can be seen in Figures 8 and 9, the quality of the approximation becomes worse as $\delta$ gets smaller. This is expected because with small number of activations, reducing their bandwidth $\delta$ leads to a worse performance as they are all zero on a large part of the space. Put differently, the number of neurons is too small to guarantee convergence of SGD to the predictions of the Wasserstein gradient flow theory.

Main results

We now state our result concerning the convergence of the SGD dynamics (3.5) to the PDE (3.9). Note that this result does not require concavity of $f$. Its proof is presented in Appendix D.

Assume that conditions (A1), (A3)-(A5) hold. Consider the SGD update (3.5) with initialization $({\boldsymbol{w}}_{i}^{0})_{i\leq N}\sim_{\rm i.i.d.}\rho_{\mbox{\tiny\rm init}}^{\delta}$ and constant step size $\varepsilon$. For $t\geq 0$, let $\rho_{t}$ be the unique solution of the PDE (3.9) with initial and boundary conditions (3.10), and assume ${\rm supp}(\rho^{\delta}_{\mbox{\tiny\rm init}})\subseteq{\sf B}({\boldsymbol{0}},r)$ Then, for any fixed $t\geq 0$, $\rho^{(N)}_{\lfloor t/\varepsilon\rfloor}\Rightarrow\rho_{t}$ almost surely along any sequence ($N,\varepsilon=\varepsilon_{N}$) such that $N\to\infty$, $\varepsilon_{N}\to 0$.

Our proof is based on the same approach developed in [MMN18]. We prove that solutions of the PDE (3.9) are in correspondence with distributions over trajectories $(\boldsymbol{X}_{t})_{t\geq 0}$ in $\Omega$ satisfying the following stochastic differential equation

where $(\boldsymbol{B}_{t})_{t\geq 0}$ is a standard Brownian motion and ${\rm d}{\boldsymbol{\Phi}}_{t}$ is the boundary reflection (in the sense of a Skorokhod problem). The density $\rho_{t}$ is determined, self consistently, via $\rho_{t}={\rm Law}(\boldsymbol{X}_{t})$. We prove existence and uniqueness of solutions to this problem, and refer to the corresponding stochastic process $(\boldsymbol{X}_{t})_{t\geq 0}$ as nonlinear dynamics. This in turn implies existence and uniqueness of the solutions of the PDE (3.9).

We next construct a coupling between the network weights $({\boldsymbol{w}}^{k}_{1},\dots,{\boldsymbol{w}}^{k}_{N})\in(\Omega^{\delta})^{N}$, and $N$ i.i.d. trajectories of the nonlinear dynamics $(\boldsymbol{X}^{t}_{1},\dots,\boldsymbol{X}^{t}_{N})\in(\Omega^{\delta})^{N}$. Controlling the expected distance in this coupling yields Theorem 5.1.

The error term in Eq. (5.1) is completely analogous to the error in a similar theorem proved in [MMN18]. The constant $\delta^{-d}$ appearing here is obtained by bounding the Lipschitz constant of $\nabla\Psi({\boldsymbol{w}};\rho)$. As already mentioned, the main technical difficulty with respect to [MMN18] is posed by the Neumann (reflecting) boundary conditions. Indeed, even if we are given a solution of the PDE (3.9), existence and uniqueness of solutions of the Skorokhod problem (5.3) is a highly non-trivial fact first established in [Tan79, LS84]. As a consequence, while the main proof idea is similar to the one in [MMN18], its implementation is significantly different.

2 Convergence to the solutions of porous medium equation

We next prove that the solution of the PDE (3.9) converges, as $\delta\to 0$, to the unique solution of the porous medium equation (3.12). As for Theorem 5.1, this result does not rely on the concavity assumption for $f$.

Assume that conditions (A1) and (A3)-(A5) hold. Denote by $\rho^{\delta}$ the unique solution of the PDE (3.9) with initial condition $\rho^{\delta}_{0}=\rho_{\mbox{\tiny\rm init}}$. Then

The porous medium equation (3.12) admits a weak solution $\rho:(t,{\boldsymbol{x}})\mapsto\rho_{t}({\boldsymbol{x}})$ with initial and boundary conditions (3.13). Further, this solution is unique under the additional condition $\rho\in\mathscrsfs{L}^{4}([0,T]\times\Omega)$.

For almost all $t\in[0,T]$, we have $\rho_{t}^{\delta}\to\rho_{t}$ in $\mathscrsfs{L}^{2}(\Omega)$ as $\delta\to 0$.

While this statement is very natural at a heuristic level, its proof is actually the bulk of our technical work. Similar approximation results have been proved in the past by Oelschläger, Philipowski, Figalli [Oel02, Phi07, FP08], but they do not apply directly to the present case unless $f=0$ (also, we have to deal with different boundary conditions).

Our proof follows a classical compactness argument, generalizing the approach of [FP08]. Namely we consider the sequence of trajectories $(\rho^{\delta}_{t})_{t\in[0,T]}$ indexed by the width $\delta$. We prove that that this family is bounded and equicontinuous in $\mathscrsfs{C}([0,T],\mathscrsfs{P}_{2}(\Omega))$, and hence admits converging subsequences $(\rho^{\delta_{n}}_{t})_{t\in[0,T]}\to(\rho_{t})_{t\in[0,T]}$. We next prove that any such converging subsequence converges in $\mathscrsfs{L}^{2}(\Omega\times[0,T])$ and that the limit is a weak solution of the porous medium equation (3.12). Unfortunately, uniqueness of weak solutions of the PME (3.12) is –to the best of our knowledge– an open problem. However, we generalize methods from [Oel02] to show that any subsequential limit is actually in $\mathscrsfs{L}^{4}(\Omega\times[0,T])$, and prove that the weak solution is unique under this condition. This allows us to conclude that $(\rho^{\delta}_{t})_{t\in[0,T]}$ converges to this unique weak solution $(\rho_{t})_{t\in[0,T]}$.

3 Global convergence of SGD

Let us now state the main result of this paper: SGD converges to a model with nearly optimal risk.

Consider the SGD update (3.5) with initialization $({\boldsymbol{w}}_{i}^{0})_{i\leq N}\sim_{\rm i.i.d.}\rho_{\mbox{\tiny\rm init}}$ and constant step size $\varepsilon$. Assume ${\rm supp}(\rho_{\mbox{\tiny\rm init}})\subseteq{\sf B}({\boldsymbol{0}};r)$. Then, for any $k\leq T/{\varepsilon}$, the following holds with probability at least $1-1/z$,

The error term $2\tau\,\Delta^{\prime}(k,{\varepsilon},d)$ in Eq. (5.4) is always non-negative. In fact, $\Delta^{\prime}(k,{\varepsilon},d)\geq 0$ as $S(\rho)\leq\log|\Omega|$ for any $\rho\in{\mathcal{P}}_{2}(\Omega)$. Furthermore, by applying Jensen’s inequality, we have that, for any $\rho\in{\mathcal{P}}_{2}(\Omega)$,

Recall that $\tau$ controls the variance of the noise, which is added at each step of the SGD algorithm for technical purposes. Thus, we can take $\tau$ sufficiently small so that the term $2\tau\Delta^{\prime}(k,{\varepsilon},d)$ is arbitrarily small.

The term $\Delta_{1}$ bounds the error due to describing the SGD dynamics using the PDE (3.9). It vanishes when $N\to\infty$, ${\varepsilon}\to 0$, under the stated conditions. The term $\Delta_{2}$ captures the error due to approximating the PDE (3.9) with the porous medium equation (3.12). Finally, the term $e^{-2\alpha k\varepsilon}$ describes the convergence to equilibrium of the solution of the porous medium equation.

The proof of Theorem 5.3 is presented in Appendix F and relies crucially on regularity results for the PDE (3.9) which are established in Appendix E.

More specifically, the proof is based on three steps, which we spell out once more:

We approximate the dynamics of SGD by the PDE (3.9) at $\delta>0$ fixed. In doing so, we incur an error $\Delta_{1}$ which is controlled using Theorem 5.1.

We approximate the solution $\rho^{\delta}_{t}$ of the PDE (3.9) at $\delta>0$ using the solution $\rho_{t}$ of the porous medium equation (3.12), as stated in Theorem 5.2.

We use results from [CJM+01, CMV03, CMV06] to prove that the latter solution converges exponentially fast to the global optimum, with rate $O(e^{-2\alpha t})$.

Given Theorems 5.1, 5.2, and the results of [CJM+01, CMV03, CMV06], this proof is relatively direct. We emphasize that, unlike Theorems 5.1, 5.2, the proof Theorem 5.3 relies in a crucial way on our structural assumptions, namely the concavity of $f$, and the structure of the bump-like activation $K_{\delta}({\boldsymbol{x}}-{\boldsymbol{w}}_{i})$.

Discussion

It is instructive to compare the general strategy followed in this paper (and in related work, e.g. [MMN18, MMM19]) and the results we obtain, to a more classical approach in theoretical statistics. For the sake of clarity, we will abstract away most of the details of the present problem, and focus on the most important differences.

This is normally proved through a uniform convergence argument to establish a bound $\sup_{{\boldsymbol{w}}}|\widehat{R}_{n}({\boldsymbol{w}})-R({\boldsymbol{w}})|\leq{\sf err}(D,n)/2$. Here ${\sf err}(D,n)$ is an error term that (hopefully) vanishes as $n\to\infty$ for $D$ fixed. Second, one proves that gradient descent (with respect to the cost function $\widehat{R}_{n}$) converges to a minimizer $\hat{\boldsymbol{w}}_{n}$. This is achieved by showing that, with high probability, the landscape ${\boldsymbol{w}}\mapsto\widehat{R}_{n}({\boldsymbol{w}})$ satisfies some strong conditions that guarantee convergence of gradient descent (or other algorithms). For instance, one desirable (although not sufficient) property is that $\widehat{R}_{n}$ does not have local minima other than the global minima, provided that the sample size is large enough. A substantial literature applies this general scheme (with significant refinements) to a variety of non-convex problems in high-dimensional statistics, including phase retrieval, clustering, matrix completion, error-in-variables models, and so on. We refer to [MBM+18] for examples and a more detailed survey.

Unfortunately this approach runs into substantial difficulties when treating complex models such as multi-layer neural networks. We can name at least two sources of difficulties. First of all, the number of parameters $D$ in the model is often comparable with the sample size $n$, and therefore uniform convergence of the empirical risk to population risk does not hold. For instance, in the present model, we could use a number of parameters $Nd\gtrsim n$: indeed, such an example is considered in Figure 5-(a), where $Nd=800$ and $n\in\{100,\dots,2000\}$. Of course this problem can be addressed by constraining other measures of complexity than the number of parameters [Bar98], but the common practice is not to add such regularizers in the training.

The second source of difficulties is that studying the risk landscape, and ruling out local minima is extremely difficult, even if we limit ourselves to the $n=\infty$ limit, i.e. the population risk $R({\boldsymbol{w}})$. In two-layers neural networks, part of this difficulty is due to the fact that the risk (1.2) is invariant under permutations of the $N$ neurons, and hence it has (generically) at least $N!$ global minima related by permutations, and a large number of saddle points connecting them.

The approach pursued in this paper builds on two simple remarks, which are connected to the previous difficulties:

Uniform convergence of the empirical risk $\widehat{R}_{n}({\boldsymbol{w}})$ to the population risk $R({\boldsymbol{w}})$ is not necessary, nor it is necessary to control the random deviations of the whole landscape of the empirical risk. What is instead important is to control the landscape of the empirical risk along the trajectory of gradient descent from a given initialization.

A convenient way to implement this idea is to consider SGD in a one-pass setting in which each sample is used only once. In the limit of small step size, this converges to gradient flow with respect to $R({\boldsymbol{w}})$.

Absence of local minima in the population landscape $R({\boldsymbol{w}})$ is not necessary either. What is instead important is absence of local minima along the gradient flow trajectory for $R({\boldsymbol{w}})$ or, more precisely, the fact that the gradient flow trajectory converges to a global minimum.

These remarks suggest the following proof strategy. Let ${\boldsymbol{w}}(t)$ denote the gradient flow trajectory from a given initialization ${\boldsymbol{w}}(0)={\boldsymbol{w}}_{0}$ (namely $\dot{{\boldsymbol{w}}}(t)=-\nabla R({\boldsymbol{w}}(t))$), and ${\boldsymbol{w}}^{k}$ be the (random) parameters produced after $k$ SGD steps. We first prove that gradient flow converges to a global optimum, possibly with explicit convergence rate $\Delta(t)$:

where $\Delta(t)\to 0$ as $t\to\infty$. We then show that the SGD trajectory, after $k$ steps, is well approximated by the gradient flow for $R({\boldsymbol{w}})$ provided the step size ${\varepsilon}$ is small. For instance we might prove that there exists a numerical constant $c_{0}$ such that, for any $k{\varepsilon}\leq T$, with high probability

The reader might recognize that the last estimate is analogous to the one obtained in Theorem 5.1, while the estimate 6.2 is what we obtain from displacement convexity (after taking the limit $\delta\to 0$ using Theorem 5.2). Putting the two estimates together, and recalling that we can run a total of $n$ SGD steps (in the one-pass setting), we get

where we set $\hat{\boldsymbol{w}}={\boldsymbol{w}}^{k}$. The error is reminiscent of a bias-variance tradeoff: the first term is a bias due to early stopping; the second is instead the stochastic approximation error. We can now optimize $n$ as to minimize this error. For instance, if $\Delta(t)=e^{-c_{1}t}$, and ${\sf err}(T)=e^{c_{2}T}$, we can choose ${\varepsilon}\propto(\log n/n)$, yielding $R(\hat{\boldsymbol{w}})\leq\min_{{\boldsymbol{w}}}R({\boldsymbol{w}})+C(\log n)^{c_{0}}/n^{c^{\prime}}$ where $c^{\prime}=c_{0}c_{1}/(c_{1}+c_{2})$.

In summary, within the present approach, the generalization error is bounded via a tradeoff between the convergence rate of gradient flow in the population risk, and the error of approximating the gradient flow by SGD. A side benefit of this proof strategy is that it guarantees the existence of an efficient algorithm to compute the weights $\hat{\boldsymbol{w}}$.

Because of these additional challenges, our bounds are not nearly as neat as in Eqs. (6.2), 6.3 and depend on the additional parameters $d,\delta$: in particular, the approximation by the porous medium equation in Theorem 5.2 is non-quantitative. We therefore refrain from optimizing the tradeoff between convergence rate of gradient flow, and error in stochastic approximation, which would result in suboptimal statistical guarantees, and defer this objective to future work.

Acknowledgements

A. Javanmard was partially supported by an Outlier Research in Business (iORB) grant from the USC Marshall School of Business, a Google Faculty Research award and the NSF CAREER award DMS-1844481. M. Mondelli was supported by an Early Postdoc.Mobility fellowship from the Swiss National Science Foundation and by the Simons Institute for the Theory of Computing. A. Montanari was partially supported by grants NSF DMS-1613091, CCF-1714305, IIS-1741162 and ONR N00014-18-1-2729. This work was carried out in part while the authors were visiting the Simons Institute for the Theory of Computing.

Appendix A Uniqueness of weak solutions of limit PDE (δ=0𝛿0\delta=0)

In this appendix, we prove that the limit PDE obtained for $\delta\to 0$, namely the porous medium equation (3.12) has at most one solution in $\mathscrsfs{L}^{4}(\Omega\times[0,T])$. Existence of such solutions will follow from the results of Appendix F, and in particular from Lemma F.4.

Throughout this appendix, we adopt the notation $\Phi(\rho)=\tau\rho+\nu_{0}\,\rho^{2}/2$. Let us formally define the concept of weak solutions for the PDE (A.1).

We say that $\rho\in\mathscrsfs{C}([0,T],$ $\mathscrsfs{P}_{2}(\Omega))$ is a weak solution of the PDE (A.1), with initial and boundary conditions (A.2) if

$\rho_{t}$ has density $\rho(\,\cdot\,,t)$ with respect to Lebesgue measure, and $\rho\in\mathscrsfs{L}^{2}(\Omega\times[0,T])$.

For any test function $h\in\mathscrsfs{C}^{2,1}(\Omega\times[0,T])$, satisfying $\langle{\boldsymbol{n}}({\boldsymbol{x}}),\nabla h({\boldsymbol{x}},t)\rangle=0$ for all ${\boldsymbol{x}}\in\partial\Omega,t\in[0,T]$, we have

We now prove a uniqueness result, under a mild integrability condition.

Note that setting $\nu_{0}=1$ corresponds to scaling time by a factor $\nu_{0}$ and to substituting $\tau$ with $\tau\,\nu_{0}$. Since the proof holds for any $\tau>0$, without loss of generality we can set $\nu_{0}=1$.

Here, $\hat{\eta}_{t}$ is a smooth approximation of $\eta_{t}^{M}$, such that $\tau\leq\hat{\eta}_{t}({\boldsymbol{x}})\leq M$. (We will make precise below in what sense $\hat{\eta}_{t}$ has to approximate $\eta_{t}^{M}$. For the moment, it can be a general smooth function satisfying the bounds $\tau\leq\hat{\eta}_{t}({\boldsymbol{x}})\leq M$.) Note that (A.6) is a backward parabolic problem with smooth coefficients and with Neumann boundary conditions. Hence, by classical results on quasilinear parabolic PDEs [LSU88], it admits a solution $h_{t}\in\mathscrsfs{C}^{2,1}(\Omega\times[0,T])$. Rewriting (A.5) for such a test function $h_{t}$, we get

By applying Cauchy-Schwarz inequality, we have that

Here $(a)$ follows from integration by parts in the integral over $\Omega$ and using the fact that $\langle{\boldsymbol{n}}({\boldsymbol{x}}),\nabla h_{t}({\boldsymbol{x}})\rangle=0$ for ${\boldsymbol{x}}\in\partial\Omega$ and $t\in[0,T]$. Also, $(b)$ follows from integration by parts in the integral over $t$. Finally $(c)$ holds because $h_{T}({\boldsymbol{x}})=0$ for ${\boldsymbol{x}}\in\Omega$ and $\mu(0)\geq 0$.

Getting back to (A) and using the properties of function $\mu(t)$, we have

The penultimate step follows from integration by parts and the constraint $\langle\nabla h_{t}({\boldsymbol{x}}),{\boldsymbol{n}}({\boldsymbol{x}})\rangle=0$, for ${\boldsymbol{x}}\in\partial\Omega$ and $t\in[0,T]$, and the last step follows by applying Cauchy-Schwartz inequality. We continue by applying Cauchy-Schwartz inequality again to get

where $C=\sup_{{\boldsymbol{x}}\in\Omega}|\nabla f({\boldsymbol{x}})|$. Combining Equations (A.11) and (A.12), we get

where $\mu_{\max}=\sup_{t\in[0,T]}\mu(t)$. We find a smooth function $\mu(t)$ such that

$\mu(t)\geq\mu_{\min}>0$, for $t\in[0,T]$,

$\mu^{\prime}(t)-\frac{2C^{2}}{\tau}\mu(t)\geq 0$.

Now by employing (A.14) in bound (A) combined with (A.7) we get

Call the first integral $I_{1}$ and denote the second one by $I_{2}$. The integrand in $I_{2}$ is pointwise bounded by

where ${\varepsilon}$ is an arbitrary small fixed constant.

In addition, since $\hat{\eta}_{t}({\boldsymbol{x}})\geq\tau$, invoking (A.15) we have

Since $\frac{\mu_{\max}}{\mu_{\min}}=e^{\frac{2C^{2}}{\tau}T}<\infty$ and $\theta$ are independent of ${\varepsilon}$, by choosing ${\varepsilon}$ arbitrarily small, we conclude that

Since $\theta_{t}({\boldsymbol{x}})\geq 0$ was an arbitrary smooth function supported on $\Omega\times[0,T]$, this implies that $u\leq 0$, almost everywhere. By repeating a similar argument, we get $u\geq 0$, almost everywhere. The result follows. ∎

Appendix B General results on the PDE (3.9) (δ>0𝛿0\delta>0)

This appendix contains some basic results on the PDE (3.9). Although these facts are standard, we collect them here for the reader’s convenience.

In fact, we will consider a more general PDE, which also includes as a special case the one studied in [MMN18]. We consider a compact convex domain $D$, with a non-empty interior. The general PDE is parametrized by two functions $V\in\mathscrsfs{C}^{2}(D)$ and $U\in\mathscrsfs{C}^{2}(D\times D)$, with $U({\boldsymbol{x}}_{1},{\boldsymbol{x}}_{2})=U({\boldsymbol{x}}_{2},{\boldsymbol{x}}_{1})$. (Unlike in [MMN18], we consider the case of a compact domain with Neumann boundary conditions.) Given $\rho\in\mathscrsfs{P}_{2}(D)$, we define

We will typically write $\rho_{t}(\,\cdot\,)$ for a solution of this equation, in order to emphasize that it is a function of $t$ that takes values in $\mathscrsfs{P}_{2}(D)$, and $\rho({\boldsymbol{x}},t)$ for the corresponding density, viewed as a function on $D\times[0,T]$. Let us formally define the concept of weak solutions for the PDE (B.2).

Note that the PDE (3.9) is a special case of this setting with $D=\Omega^{\delta}$, and $V({\boldsymbol{w}})$ and $U({\boldsymbol{w}}_{1},{\boldsymbol{w}}_{2})=U({\boldsymbol{w}}_{1}-{\boldsymbol{w}}_{2})$ defined as follows:

For the special choice of $V$ and $U$ given by (B.4) the following properties hold:

$\lim_{\delta\to 0}\sup_{{\boldsymbol{w}}\in\Omega^{\delta}}|V({\boldsymbol{w}})+\nu_{0}\,f({\boldsymbol{w}})|=0$.

$U({\boldsymbol{w}})=\nu_{0}\,\delta^{-2d}K^{(2)}({\boldsymbol{w}}/\delta)$, where $K^{(2)}=K*K$.

We have $V({\boldsymbol{w}})=-\nu_{0}\int K^{\delta}({\boldsymbol{x}})f({\boldsymbol{w}}-{\boldsymbol{x}}){\rm d}{\boldsymbol{x}}$. Hence,

This proves that $V({\boldsymbol{w}})$ is convex. The next two properties are straightforward. ∎

We say that $\rho:[0,T]\to\mathscrsfs{P}_{2}(D)$ is a weak solution of (B.2) with initial and boundary conditions (B.3) if $\rho\in\mathscrsfs{C}([0,T],\mathscrsfs{P}_{2}(D))$ and, for any test function $h\in\mathscrsfs{C}^{2,1}(D\times[0,T])$, satisfying $\langle{\boldsymbol{n}}({\boldsymbol{x}}),\nabla h({\boldsymbol{x}},t)\rangle=0$ for all ${\boldsymbol{x}}\in\partial D,t\in[0,T]$, we have

We now state and prove Duhamel’s principle for the PDE (B.2). Duhamel’s principle follows from the fact that the right-hand side of (B.2) contains the linear diffusion term $\tau\Delta\rho$, and it will be crucial for the proofs that will follow.

Assume $\tau>0$. Let $G^{D}({\boldsymbol{x}},{\boldsymbol{y}};t)$ denote the heat kernel with Neumann boundary conditions, defined in (G.1)-(G.3). Let $\rho$ be a weak solution of the PDE (B.2) with initial and boundary conditions (B.3). Then, for any $t>0$, $\rho_{t}({\rm d}{\boldsymbol{x}})$ has a density, denoted by $\rho(\;\cdot\;,t)$, which satisfies, for any $t>0$,

By rescaling time, without loss of generality, we set $\tau=1$. Let $\varphi\in\mathscrsfs{C}^{2}(D)$, and define

By the properties of the heat kernel, we have:

Let $\rho_{t}$ be a weak solution. We choose the test function $h({\boldsymbol{x}},s)=G_{\varphi}({\boldsymbol{x}};t-s)$ in (B.5) with $T=t$. Note that by (B.8), this test function satisfies the Neumann boundary condition. In addition, by (B.9) we obtain

By an application of Fubini’s theorem, this implies

Since $\varphi\in\mathscrsfs{C}^{2}(D)$ is arbitrary, we obtain that $\rho_{t}$ admits a density and (B.6) follows. ∎

As an intermediate step towards proving existence and uniqueness, we consider a linearized problem

Then, the PDE (B.13) with initial and boundary conditions (B.14) has at most one weak solution.

Without loss of generality, we will set $\tau=1$. Assume by contradiction that $\rho^{(1)}$, $\rho^{(2)}$ are two solutions. Fix arbitrary $0\leq t^{\prime}\leq t$. Then, by an application of (B.6) to $\Psi_{*}({\boldsymbol{x}},t)$, we have

where we used the estimates of Theorem G.1. By taking supremum over $0\leq t^{\prime}\leq t$ form both sides, we obtain that for $t<1/(C(D)^{2}\|\nabla\Psi_{*}\|_{\mathscrsfs{L}^{\infty}(D\times[0,T])}^{2})$,

Therefore, the two solutions coincide if we fix the initial condition $\rho^{(1)}(\,\cdot\,,0)=\rho^{(2)}(\,\cdot\,,0)=\rho_{\mbox{\tiny\rm init}}$. For larger $t$, the claim follows by iterating the above argument. ∎

Appendix C Nonlinear dynamics

The ‘nonlinear dynamics’ plays an important role in our proof of Theorem 5.1. In this section we adopt the same general setting as in Appendix B, remembering that for our application we set $D=\Omega^{\delta}$ and $U,V$ as per Eq. (B.4).

Given $\rho:[0,T]\to\mathscrsfs{P}_{2}(D)$, consider the following stochastic differential equation for a process $(\boldsymbol{X}_{t})_{t\in[0,T]}$, with a reflecting boundary condition (known as ‘Skorokhod problem’)

where $(\boldsymbol{B}_{t})_{t\geq 0}$ is a standard $d$-dimensional Brownian motion and $({\boldsymbol{\Phi}}_{t})_{t\geq 0}$ enforces the reflecting boundary by satisfying the following constraints (recall that ${\boldsymbol{n}}({\boldsymbol{x}})$ is the normal to $\partial D$ at ${\boldsymbol{x}}\in\partial D$, directed inside):

$({\boldsymbol{\Phi}}_{t})_{t\geq 0}$ is adapted (and hence so is $(\boldsymbol{X}_{t})_{t\geq 0}$).

$t\mapsto{\boldsymbol{\Phi}}_{t}$ has (almost surely) bounded variation. Denoting by $\|{\boldsymbol{\Phi}}\|_{\mbox{\tiny\rm TV}}(t)$ the total variation of ${\boldsymbol{\Phi}}$ on the interval $[0,t]$, we define the measure $\mu_{\Phi}$ on $[0,T]$ by $\mu_{\Phi}([0,t])=\|{\boldsymbol{\Phi}}\|_{\mbox{\tiny\rm TV}}(t)$.

$\mu_{\Phi}(\{t:\,\boldsymbol{X}_{t}\in D^{\circ}\})=0$, where $D^{\circ}$ denotes the interior of $D$.

where ${\boldsymbol{N}}_{s}={\boldsymbol{n}}(\boldsymbol{X}_{s})$, for $\mu_{\Phi}$-almost every $s$.