On the Convergence of Gradient Descent Training for Two-layer ReLU-networks in the Mean Field Regime

Introduction

Practitioners have found that artificial neural networks can be trained by gradient descent-based algorithms to fit many complicated target functions. While it is well-understood why the function class is sufficiently expressive for this purpose [Bar93, Cyb89, Hor91, LLPS93], the choice of optimal network parameters in applications is a highly non-convex problem. It is not fully understood why cleverly initialized gradient descent achieves good performance in practice.

In [CB18b], the authors prove that if the parameter distributions of neural network-like models converge to a limiting distribution, then the limit is in fact a global minimizer. The result is true for a class of ‘spread out’ initial conditions (implying very wide networks) and under some technical assumptions. One of the most prominent cases concerns function models with the same homogeneity as neural networks with a single hidden layer and ReLU (or leaky ReLU) activation. However, the result does not apply directly due to the lack of differentiability of these activation functions.

In this article, we extend the main result of [CB18b] for this ReLU-like setting and generalize previous results for some cases. The results proved here improve upon previous work in two ways:

Our analysis applies to ReLU-networks with suitable initial conditions rather than toy models with similar properties.

We only assume that a limiting object, whose existence is guaranteed by compactness, is unique.

In the analysis we exploit that ReLU activation

Consider a two-layer mean field network model

As time goes to infinity, $\mathcal{R}(\pi_{t})$ converges to minimum Bayes risk.

The velocity potential $\frac{\delta\mathcal{R}}{\delta\pi}$ converges to a unique limit $g$ locally uniformly as $t\to\infty$.

If $\mathcal{R}(\pi_{t})$ approaches MBR, we can additionally identify the limit

The study of Wasserstein gradient flows in the context of (infinitely wide) shallow neural networks is motivated below. The same problem exactly captures the gradient descent training of finite neural networks for infinitesimal learning rate and asymptotically captures stochastic gradient descent when the learning rate is so small/batch size is so large that stochastic effects are negligible.

As time goes to infinity, $\mathcal{R}(\pi_{t})$ converges to minimum Bayes risk.

The velocity potential $\frac{\delta\mathcal{R}}{\delta\pi}$ converges to locally uniformly as $t\to\infty$.

The loss function is assumed to be $C^{1}$-smooth (with bounded derivative). We can consider more general loss functions (e.g. mean squared error), but only for uniformly bounded data. The initial condition can be generalized, but we require that at time $t=0$ we have $-a^{2}+|w|^{2}+b^{2}\geq 0$ $\pi$-almost surely. The Minkowski inner product is preserved along the gradient evolution and the condition $(w,b)\neq 0$ unless also $a=0$ allows us to avoid issues of non-differentiability. To establish existence of a minimizer, it suffices that a certain projection of $\pi_{t}$ converges.

The significance of our result is as follows.

We do not need to assume the existence of a limit (guaranteed by compactness), but only its uniqueness.

One of the main complications of the question of convergence to minimal energy for this gradient flow is the dependence on the initial condition. The question whether the limit of the velocity potential is generically unique can be asked independently of the initial condition and seems more approachable by standard means of analysis. We therefore believe that this perspective might be a first step towards the convergence theory for Wasserstein gradient flows for shallow neural networks.

In the proof, we show directly that convergence to minimal energy implies convergence of the velocity potentials $\frac{\delta\mathcal{R}}{\delta\pi}$ to zero. This resembles the first order optimality condition in classic calculus and does not require deep geometric insight. In the converse direction, the geometry of the energy landscape is used via the homogeneity of the activation function. While we formulate all results for ReLU-activation, they hold equivalently for leaky ReLU networks. We prove that the second moments of the evolving parameter distribution $\pi_{t}$ grow at most sublinearly under general conditions and show that if the unique limit $g$ does not vanish identically, they grow exponentially. Thus only $g\equiv 0$ is admissible as a unique limit. In this situation the second moments – which are bounded below by zero – decrease linearly at a positive rate unless $\mathcal{R}(\pi_{t})$ decays to minimum Bayes risk.

The article is structured as follows. In the remainder of this section, we briefly review previous work and collect some notation. In Section 2, we describe the setting, state our main assumptions, and review some additional background information on Wasserstein gradient flows and continuity equations. In Section 3, we state our main results. Some concrete examples of data distributions and loss functionals to which the results apply are listed in Section 4. We conclude with a discussion of our results in Section 5. Longer proofs and some technical details are collected in Appendix A. Appendix B is dedicated to the technical condition of Morse-Sard type which we assume.

Introductions to machine learning in general and neural networks in specific can be found for example in [SSBD14, HH19, GBC16, MBW+19].

The question why (stochastic) gradient descent starting at a suitable random initialization finds good parameters for artificial neural networks despite the fact that the energy landscape is highly non-convex has attracted much attention and several competing explanations have emerged.

One avenue of research aims to uncover a description of gradient descent in neural networks in the infinite neuron limit in the mean field scaling regime [CB18b, RVE18, SS20, MMN18]. Convergence criteria in the shallow network setting are developed for example in [AKSG19] and [CB18b], but have not been established generally.

Another line of articles [EMWW19, EMW19d, BM19, DZPS18, DLL+18, JGH18, ADH+19] considers a heavily overparametrized regime with large initialization. In this setting, the gradient flow for neural networks behaves is proved to behave like the gradient flow of a very wide random feature model [EMW19d, EMWW19] with high probability over the choice of (suitable) initial condition. In particular, the direction and bias of a neuron barely change from their (random) initialization in this regime. Chizat and Bach dub this the ‘lazy training regime’ since neurons hardly move. They show that the underlying analysis is due rather to a (usually implicit) scaling assumption on initialization rather than the specific structure of neural networks [CB18a].

The success of neural networks in practical applications has been explained by the observation that – unlike any linear theory – neural networks can beat the ‘curse of dimensionality’ [Bar93]. It thus seems unlikely that linearization is able to explain their recent success. Furthermore, in these studies parameters are usually initialized so large that the path norm a two-layer network with $m$ hidden neurons scales like $\sqrt{m}$ at initialization. Natural generalization bounds as derived in [EMW19a, EMW18] therefore do not apply.

Practitioners tend to train neural networks using stochastic gradient descent rather than full gradient descent. For small learning rate (time step size), the evolution can be described by SDEs with Gaussian noise [HLLL19, LTE15]. The noise coefficient is given by the covariance of the gradient, which is usually neither isotropic nor homogeneous. However, in the mean field regime and assuming that the noise is standard Gaussian, one can prove that the parameter distribution converges to a good value (minimizer of a regularized risk functional), see [HRSS19]. The derivation is built on the link of the heat equation to both stochastic analysis and optimal transport theory [JKO98]. In this case, the parameter distribution approaches the stationary measure of a Markov process as time approaches infinity. If the noise is sufficiently large, the convergence is exponential, while for small noise, the stationary measure approaches a minimizer of the mean field risk functional.

Rigorous convergence results in realistic settings can be obtained under strong assumptions on the initial condition (or a state which arises along the gradient flow). One example is [BJ18], where the authors show that if a neural network has well chosen parameters at some time, the parameters improve along the gradient flow. In [AKSG19] the authors study a mean field gradient flow and show that under a regularity/closeness assumption, the gradient flow converges. Unfortunately, the condition cannot be verified in practice. Global convergence for a toy model with similar properties is established in [EMW19c, Section 7].

Some results in [CB18b] also apply to neural networks with smooth activation and more than one hidden layer. However, the imposition of a linear structure results in network-like models where each neuron in the outermost layer has its own trainable weights for the deeper layers. More recent works in the mean field setting consider deep neural networks whose parameters are initialized independently across the layers [AOY19, NP20, SS19]. The independence is preserved through time (‘propagation of chaos’) and a mean field description is available in different scaling limits. The theory is entirely different from that of shallow networks. While a shallow networks can be described by indexed particles $(a_{i},w_{i},b_{i})$, the paths in a deep network through multiple layers have a more complicated interacting multi-index structure $(a_{i},b_{ij},c_{j})$.

2. Notations and Terminology

Background and Assumptions

In this section, we describe the objects under consideration in this article. Some assumptions will be relaxed in Section 3.4.

For this introduction, the choice of cone $\Theta$ is inessential and will be motivated later in Section 2.5. Note that the family

The property is known to hold when the parameters $\theta_{i}$ are not constrained to a cone, the factor $\frac{1}{m}$ is not present and coefficients $a_{i}$ are included before $\phi(\theta_{i},\cdot)$ [Cyb89]. None of these differences are significant as

A parameter distribution $\pi$ is a Borel probability measure on $\overline{\Theta}$ with finite second moments

i.e. $\pi\in\mathcal{P}:=\mathcal{P}_{2}(\overline{\Theta})$ is an element of Wasserstein space over the cone $\overline{\Theta}$. We denote the realization of $\pi$ as

The uniform distribution on the unit sphere satisfies (P4).

which encodes many important properties of the problem.

2. The Risk Functional

Combining all previous notions, we define the risk functional $\mathcal{R}:\mathcal{P}_{2}(\overline{\Theta})\to[0,\infty]$,

A few observations are in order. Recall that $\mathcal{P}$ or $\mathcal{P}_{2}$ denotes the space of Radon measures with finite second moments (on the natural space in a given context, which usually is $\overline{\Theta}$).

The function $(x,\alpha)\to L_{x}(\alpha)$ is measurable (see Lemma 2.3 below), but cannot generally be assumed to be continuous.

We note that there exists a measurable selection of minimizer for $L_{x}$. This is not entirely immediate even when $L_{x}$ is strictly convex.

The functional $\mathcal{R}$ admits a minimizer if and only if there exists a measure $\pi\in\mathcal{P}$ such that

In particular, as $f_{\pi}$ is $N(\pi)$-Lipschitz (where again $N(\pi)$ denotes the second moments of $\pi$), there must be a version of $f^{*}$ which is Lipschitz-continuous. The Lipschitz condition is far from sufficient [EW20b]. Since

the set of minimizers is a convex subset of $\mathcal{P}$.

Write $T(a,w,b)=(-a,w,b)$. The fact that $\phi(T\theta,\cdot)\equiv-\phi(\theta,\cdot)$ implies that the measure $\pi$ representing a function $f_{\pi}$ cannot be unique: Any measure $\pi$ such that $T_{\sharp}\pi=\pi$ (i.e. $\pi(T^{-1}(V))=\pi(V)$ for all measurable $V\subseteq\Theta$) represents the function since

In ReLU networks, another source of non-uniqueness is the identity

More generally, if $\pi$ represents $f_{\pi}$ and $\pi^{\prime}$ represents , then for any $\lambda\in(0,1)$ consider the probability measure

In particular, while the minimizer $f^{*}$ of $\widetilde{\mathcal{R}}$ is unique if $L_{x}$ is strictly convex, a minimizer of $\mathcal{R}$ must be highly non-unique. This is a major obstacle in linearization-based approaches to convergence.

3. Wasserstein Gradient Flows

The study of Wasserstein gradient flows in machine learning is motivated by the following observation: The parameters $\{\theta_{i}\}_{i=1}^{m}\in\overline{\Theta}^{m}$ of a parametrized function

evolve by the time-accelerated Euclidean gradient flow

if and only if their distribution $\pi_{m}=\frac{1}{m}\sum_{i=1}^{m}\delta_{\theta_{i}}$ follows the $2$-Wasserstein gradient flow of the extended risk functional

A key observation is that the individual particles $\theta_{i}$ are irrelevant and only their distribution matters when computing $f_{\theta_{1},\dots,\theta_{m}}$. We refer to two-layer network functions of the form $f(x)=\frac{1}{m}\sum_{i=1}^{m}a_{i}\,\sigma(w_{i}^{T}x+b_{i})$ as mean field networks in contrast to classical two-layer networks $f(x)=\sum_{i=1}^{m}a_{i}\,\sigma(w_{i}^{T}x+b_{i})$. Both classes are identical from the perspective of approximation theory, but lead to different dynamic models in the infinite-width limit. Classical networks are described by the linearized dynamics of neural tangent kernels, while mean field networks evolve truly non-linearly by Wasserstein gradient flows.

Thus optimizing mean field network parameters by the gradient flow of a risk functional is equivalent to optimizing their distribution by a Wasserstein gradient flow, see e.g. [CB18b, Proposition B.1]. An expanded heuristic can also be found in Appendix A.

The gradient flow and the map $\pi\mapsto f_{\pi}$ (so also the risk functional $\mathcal{R}$) are naturally defined on the Wasserstein space $\mathcal{P}_{2}$ of probability measures with finite second moments. Background on optimal transport theory and Wasserstein gradient flows can be found e.g. in [AGS08, San15, Vil08].

In this article, we mostly consider Wasserstein gradient flows of continuous distributions. By continuity, the Wasserstein gradient flows $\pi_{m}(t)$ starting at $\pi_{m}^{0}$ converge to the solution of the Wasserstein gradient flow $\pi(t)$ starting at $\pi^{0}=\lim_{m\to\infty}\pi_{m}^{0}$ for all $t>0$. Even more, the limits

commute (if they exist). A proof under an additional technical condition (which can be eliminated if one considers gradient flows of unregularized risk or the second moment regularizer) is given in Appendix B of [CB18b]. The complication arises from regularizing functionals which do not have the correct homogeneity, which leads the authors to consider compactly supported initial conditions.

Hence if the gradient flow starting at a measure $\pi^{0}$ converges to a minimizer of risk, then gradient flows starting at closeby empirical measures asymptotically achieve low risk. The theory, at this point, is purely qualitative, but nonetheless indicative for practical applications.

The Wasserstein gradient flow $\pi_{t}$ is described by the continuity equation

denote the variational derivative of $\mathcal{R}$ with respect to $\pi$ and its spatial gradient respectively.

4. Continuity Equations

Then according to [Amb08, Proposition 4], we have

The Lemma does not apply directly to the situation which we will consider since the flow field $V$ will be positively one-homogeneous and thus unbounded in most cases. However, the result also applies under the weaker assumption

see [Amb08, Remark 7]. The condition of at most linear growth is required to prevent particles from escaping to infinity in finite time. Thus (2.5) and (2.6) apply also in our situation.

5. Initial Condition

There are two considerations concerning the initial parameter distribution of the gradient flow. The first is specific to ReLU activation and of a technical nature while the second one is more geometric and concerns energy decay to minimum Bayes risk/convergence to minimizers.

We can formally compute the parameter gradient of the activation function

There is a more subtle problem of regularity. Assumption (P4) guarantees that

generally fails to even be continuous. This difficulty can only be circumvented if $a$ is close to zero whenever $(w,b)$ is close to zero. At this point, the geometry of the ReLU function as the product of two positively one-homogeneous functions comes into play. Namely, we can use the fact that

In particular, the cone of space-like vectors

is preserved along the gradient flow evolution. Inside $\overline{\Theta}$, we have $|a|\leq\sqrt{|w|^{2}+b^{2}}$, which allows us to save the Lipschitz property. We therefore restate an assumption on the initial condition more explicitly.

It is possible to consider initial conditions supported on other super-level sets like $\Theta_{\varepsilon}=\{(a,w,b):-a^{2}+|w|^{2}+b^{2}>\varepsilon^{2}\}$ instead. A flow starting at $\pi_{0}$ supported on $\Theta_{\varepsilon}$ will never reach a point where $(w,b)=0$, but on the other hand does not allow us to exploit the positive two-homogeneity of $\phi$ in $\theta$ as easily since $\Theta_{\varepsilon}$ is not a cone. Positive two-homogeneity is essential to many arguments below, so we proceed with $\Theta=\Theta_{0}$ instead.

5.2. Omni-directional initial conditions

It is well-known that the functional $\mathcal{R}$ is not sufficiently convex in Wasserstein geometry to guarantee convergence to global minimizers from any initial condition. In particular, if a global minimizer $\pi$ exists but cannot be written as an empirical measure with $m$ atoms, any initial condition corresponding to an empirical measure with $m$ atoms cannot converge to $\pi$ since the continuity equation has no smoothing effect and preserves atomic measures.

The following class of initial conditions is successful in theory and applications.

We call a probability measure $\pi$ on $\Theta$ omni-directional if every open cone in $\Theta$ has positive measure.

Omnidirectional initial condition: $\pi_{0}$ is omni-directional.

Clearly, an omni-directional initial condition is an abstraction available only in the infinite-width limit.

6. Morse-Sard Property

Below, we prove existence and uniqueness for gradient flow training of ReLU-activated two-layer networks under the assumptions listed above. For technical reasons, we add an assumption which has only been established in full generality in dimension $d=2$. The assumption is only required when discussing the limiting behavior of the gradient flow.

locally uniformly on $\overline{\Theta}$. Then the restriction of $g$ to the unit sphere $S^{d+1}=\{a^{2}+|w|^{2}+b^{2}=1\}$ has the property that $P^{S^{d}}_{*}V=\nabla^{S}g$ does not vanish anywhere on the level set $\{g=t\}$ for Lebesgue-almost every $t$, where $\nabla^{S}$ denotes the gradient of $g$ tangent to the sphere and $P^{S^{d}}_{*}$ the tangent map to the projection onto the sphere.

We discuss the Morse-Sard condition below in Appendix B. Using smoothness, we establish the property in dimension two and provide evidence on the other hand that (M-S) may not be expected to hold in high dimension in full generality. We isolate the point in the proof where the condition is used and where an argument would need to be adapted in order to avoid it.

A weaker version of the main result holds without this condition.

Evolution of the Parameter Distribution

We first establish that solutions to gradient flow training exist. Assume all conditions outlined above except for (IC2) and (M-S), which are only required for statements about limiting objects.

Let $\pi_{0}\in\mathcal{P}_{2}(\overline{\Theta})$. Then there exists a unique solution to the Wasserstein gradient flow (2.3).

The next statement will be useful to understand the behavior of limits as $t\to\infty$.

If $\pi_{0}\in\mathcal{P}_{2}(\overline{\Theta})$ is omni-directional and $\pi_{t}$ denotes the Wasserstein gradient flow starting at $\pi_{0}$, then $\pi_{t}$ is omni-directional for all $t>0$.

This Lemma is a simpler version of [CB18b, Theorem 3.3]. The proof is based on the flow map representation. Since the activation function is positively two-homogeneous in the network parameters, the flow field $-\nabla(\delta_{\pi}\mathcal{R})$ is positively one-homogeneous, which means that half-rays move as half-rays and cones are preserved under the flow. Both results also apply directly to explicitly regularized risk functions with suitable homogeneity such as

2. Growth of Second Moments

Recall that we had denoted the second moment of $\pi$ by

Under gradient flow training, $N$ can only grow sublinearly in time.

If $\pi_{t}$ evolves by the Wasserstein-gradient flow of $\mathcal{R}$, then

On the other hand, the key insight of [CB18b] is that the homogeneity of $\phi$ implies that if the velocity potentials $\delta_{\pi}\mathcal{R}$ were to converge to a non-trivial limit as $t\to\infty$, it would lead to exponential growth of $N(\pi_{t})$.

converge to a function $g$ and a vector field $V=\nabla g$ respectively, locally uniformly on $\overline{\Theta}$. If $g\not\equiv 0$, then

We give both proofs in the appendix for the reader’s convenience. The proof of Lemma 3.4 is the only point in the document where the assumptions (IC2) and (M-S) are used.

3. Main Results: Lipschitz Loss

We can now characterize the convergence of Wasserstein gradient flows for the risk functional $\mathcal{R}$ with omni-directional initial conditions entirely. We consider the simultaneous $\omega$-limit set

Assume the above conditions (L1), (L2), (L3), (P1), (P2), (P3), (P4), (LP1), (LP2), (IC1), (IC2), (M-S). Then

$\lim_{t\to\infty}\mathcal{R}(\pi_{t})=\inf_{\pi}\mathcal{R}(\pi)$ if and only if $\Omega^{lim}$ consists of only one element.

If $\Omega^{lim}$ consists of only one element, then that element is $\lim_{t\to\infty}(\delta_{\pi}\mathcal{R})(\pi_{t};\theta)\equiv 0$.

The first statement of the Theorem does not require (IC2) and (M-S). Verifying that there exists only a single element in $\Omega^{lim}$ is non-trivial. Generally uniqueness is proved by showing that the limit satisfies an equation which can be studied separately. Even assuming that there exists a limiting measure $\pi_{\infty}$ such that $\pi_{t}\to\pi_{\infty}$, the natural ‘zero dissipation in the limit’ condition

only shows that $g(\theta)=0$ for $\pi_{\infty}$-almost all $\theta$. This is significantly weaker than every $\theta$ if $\pi_{\infty}$ concentrates on a small set. Functions satisfying the zero-dissipation property are stationary points of the flow, and there are many of them which are not the global minimum. Not assuming the existence of a limit $\pi_{\infty}$, linearization becomes highly involved since the map $\pi\mapsto f_{\pi}$ is far from injective. This makes the question of uniqueness non-trivial. Despite these obstacles, we believe the result to be a first step towards a general convergence theory. Without assumptions (IC2), (M-S), a weaker result holds.

Assume that (L1), (L2), (L3), (P1), (P2), (P3), (P4), (LP1), (LP2) and (IC1) hold. Then

$\lim_{t\to\infty}\mathcal{R}(\pi_{t})=\inf_{\pi}\mathcal{R}(\pi)$ if and only if $|\Omega^{lim}|$ contains only the function $g=0$.

Having found a characterization for risk converging to zero for omni-directional initial conditions, we now turn our attention to the question of whether we can find a minimizer of risk. The corollary is almost immediate.

If in addition to the assumptions of Theorem 3.5 we assume that there exist a sequence of times $t_{n}\to\infty$ and a measure $\pi_{\infty}$ on $\Theta$ such that $\pi_{t_{n}}\to\pi_{\infty}$ in 2-Wasserstein distance and $|\Omega^{lim}|=1$, then $f_{\pi_{\infty}}=f^{*}$, i.e. $\pi_{\infty}$ is a risk minimizing measure.

The existence of a suitable subsequence is guaranteed if the $p$-th moments of $\pi_{t}$ remain uniformly bounded for some $p>2$. A version of Corollary 3.7 also holds under the same assumptions ad Theorem 3.6. Under a weaker condition, we obtain a weaker statement.

If in addition to the assumptions of Theorem 3.5 we assume that the second moments $N(\pi_{t})$ remain uniformly bounded (at least along a subsequence of times $t_{n}\to\infty$) and $|\Omega^{lim}|=1$, there exists a measure $\pi_{\infty}\in\mathcal{P}$ such that $f_{\pi_{\infty}}=f^{*}$ (i.e. a risk minimizing measure).

Again, there is a version of Corollary 3.8 under the conditions of Theorem 3.6. The measures $\pi_{t_{n}}$ may not have a limit in 2-Wasserstein distance, but we can explicitly construct $\pi_{\infty}$ from $\pi_{t_{n}}$ by a reparametrization argument. If the second moments of $\pi_{t}$ remain uniformly bounded in time, there exists a subsequence $\pi_{t_{n}}$ which converges to a limit in $p$-Wasserstein distance for all $p<2$. This is not sufficient since the map $\pi\mapsto f_{\pi}$ is not continuous in this topology, as the following example shows:

4. Main Results: Smooth Loss and Bounded Data

for any $S<\infty$. We consider the modified loss function

for $S>0$ and the modified risk functional

$\mathcal{R}_{S}(\pi)=\mathcal{R}(\pi)$ for all $\pi$ such that $N(\pi)\leq\frac{S}{1+R}$.

In particular, the gradient flow $\pi_{t}^{S}$ of $\mathcal{R}_{S}$ starting at a fixed measure $\pi_{0}$ independent of $S$ exists and $\pi_{t}^{S}=\pi_{t}^{S^{\prime}}$ describes the gradient flow of $\mathcal{R}$ for $t<{c\,\min\{S,S^{\prime}\}}/{(1+R)}$. Note that the growth rate of $N$ does not depend on $S$. Thus we may take $S$ to infinity to find that the gradient flow $\pi_{t}$ of $\mathcal{R}$ exists on $[0,\infty)$. We have shown the following.

Replace assumption (P2) by the stronger assumption

and condition (L2) by the weaker assumption that

Then the gradient flow $\pi_{t}$ of $\mathcal{R}$ starting at $\pi_{0}\in\mathcal{P}_{2}(\overline{\Theta})$ exists. If $\pi_{0}$ is omni-directional, so is $\pi_{t}$ for all $t>0$.

Also the results on the convergence of gradient flows generalize under slightly stronger assumptions on the loss function. Examining the proof of Theorem 3.5, we find that the relevant properties of the proof are the following, which we postulate as conditions for smooth loss functions.

$\inf_{\pi}\mathcal{R}(\pi)=\widetilde{\mathcal{R}}(f^{*})$.

Assume the conditions (L1), (L2’), (L3), (P1), (P2’), (P3), (P4), (LP1), (LP2), (IC1), (IC2), (M-S), (SL1), (SL2) and (SL3). Then

$\lim_{t\to\infty}\mathcal{R}(\pi_{t})=\inf_{\pi}\mathcal{R}(\pi)$ if and only if $\Omega^{lim}$ consists of only one element.

If $\Omega^{lim}$ consists of only one element, then that element is $\lim_{t\to\infty}(\delta_{\pi}\mathcal{R})(\pi_{t};\theta)\equiv 0$.

Corollaries 3.7 and 3.7 remain true also in this case, and Theorem 3.6 generalizes in the same way.

Examples

We imposed a number of abstract conditions on the loss function, data distribution, and initial condition. In this section we consider concrete situations where the conditions are met.

then any $f^{*}(x)\in$ is an admissible minimizer. More generally, we can observe that Huber loss has a characteristic smoothing length scale. If data is very spotty on a larger length scale, uniqueness of the minimum may not hold.

is almost surely not finite-valued. Our results apply directly in the first situation, but not the second. In this case, the minimizer is simple enough to be understood directly. Namely,

would be expected to lead to similar behavior. Here $\eta$ is any compactly supported probability density, $\varepsilon>0$ is a small parameter and $\mathcal{L}$ denotes Lebesgue measure. Empirical measures are recovered in the singular limit $\varepsilon\to 0$.

Using Theorem 2.1, many data distributions of practical importance are admissible for Lipschitz loss, including all distributions which have a bounded density with respect to Lebesgue measure and decay suitably fast at infinity, e.g.

distributions with continuous density on a bounded open set or

Gaussian mixture models with uniformly bounded means and variances.

For general loss functions, the first class is admissible, while Gaussian mixture models are excluded for purely technical reasons.

The question which lower-dimensional objects have admissible geometry remains open at this point.

Initial conditions need to be omni-directional and satisfy the scaling property

The easiest choice of admissible condition is to independently choose $(w,b)$ uniformly distributed on $S^{d}$ and $a$ uniformly distributed on $$. In applications, a popular choice is to initialize$(w,b)$ according to a standard normal distribution with mean zero and unit variance. It is easily computed that

and due to [Ver18, Theorem 3.1.1], the norm of $(w,b)$ concentrates close to $\sqrt{d}$ in the sense that

for a universal constant $c>0$. Thus if $a$ is distributed on a domain sufficiently smaller than $O(\sqrt{d})$ and $d$ is reasonably large with respect to the network width, then with high probability $\pi_{0}$ satisfies (IC1).

Conclusion

We have shown that the convergence of gradient descent training with infinitesimal step size for two-layer networks with ReLU or leaky ReLU activation starting at omni-directional initial conditions is equivalent to the convergence of the velocity potential to a unique limit (under certain technical conditions). The result holds for a fairly general class of loss functions and data distributions. Convergence along subsequences is guaranteed by compactness.

We have shown that convergence to minimal Bayes risk is equivalent to the convergence of the velocity potentials $\delta_{\pi}\mathcal{R}$ to a unique limit (for suitable initial conditions). Whether the limiting potential is generally unique remains one of the most relevant open questions in theoretical machine learning.

To prove existence and make use of the flow map representation, we are restricted to population risk for suitable data distributions. Especially for the case of low-dimensional data in high-dimensional spaces, the regularity condition on the data manifold is hard to understand and check.

Even if a risk-minimizing measure $\pi$ exists and risk decays to its minimum, it is unclear whether

the second moments $N(\pi_{t})$ remain bounded, and

the measures $\pi_{t}$ converge to a minimizer weakly or in $2$-Wasserstein distance (assuming that $N(\pi_{t})$ remains bounded).

State-of-the-art neural network architectures can have hundreds or even thousands of layers, far from the two-layer situation considered here. In [CB18b], the authors consider also the case of smooth bounded activation functions which are linear in one direction (and sufficiently smooth). These results apply to network architectures in which every node in the outermost layer is given its own set of parameters for deeper layers. Other models for mean field training of deep networks [AOY19, Ngu19, NP20, SS19] are very different from models for shallow networks. No analogous result is available in this setting.

Appendix A Proofs

In this section we collect the previously omitted proofs.

Thus $H$ is represented by as $H=\{x|w^{T}x+b>0\}$ for

In particular, there is a canonical representative in the unit sphere

For the derivative estimates below, we can assume without loss of generality that $b\notin\{0,\pm 1\}$, as the Lipschitz estimate extends to the boundary points by uniform continuity. Given a density $\rho$, denote

We prove the first claim. Let $(w,b)\in S^{d}$. Without loss of generality, $w=\sqrt{1-b^{2}}\,e_{1}$. Here, we can even take $\varepsilon=0$ in the decay condition and compute

This is bounded from above uniformly when $b$ is bounded away from $\pm 1$ by (A.1) and close to $b=\pm 1$ by (A.2). The passage to the limit can be justified using Fatou’s lemma. A similar estimate holds for the limit $h\to 0^{-}$, and the derivative in direction of increasing/decreasing $w$ radially is the same as that in direction of changing $b$.

Then we can therefore bound the tangential derivative as follows.

The remainder of the proof proceeds as above, except for the weight of $|x_{2}|$ in front of the density. While previously a decay as $|x|^{-(d+1+\varepsilon)}$ for $\varepsilon>0$ would have been sufficient, here we need decay as $|x|^{-(d+2+\varepsilon)}$ at infinity.

The proof of the second statement is similar to the first. The third statement is immediate from the definition. To consider the fourth statement, let ${\mathcal{M}}$ be the set of Radon probability measures

for all open sets $U$. In particular, the symmetric differences of open sets are open. This establishes the Lipschitz condition by the previous analysis since

In particular, sums of Gaussian distributions or compactly supported regular distributions as they occur in density estimation are admissible data distributions for our purposes. They can be computed from a given finite data sample and mollify the problem sufficiently for our convergence result.

Many full-dimensional data distributions satisfy (P4), but distributions with a bounded density on the hypersphere is admissible. If data is concentrated on a manifold of dimension $k$, the intuition is that $k$ should not have any ‘straight’ segments which lie mostly in a lower-dimensional affine subspace.

A complete characterization of admissible measures is beyond the scope of this article. We move on to the proof of the measurability of $f^{*}$.

Hence the function $(x,\alpha)\mapsto L_{x}(\alpha)$ is a Caratheodory integrand (finite, continuous in $\alpha$ and measurable in $x$) where $\alpha$ belongs to a second countable complete space. It is well-known that such functions are jointly measurable, see [AB94, Theorem 14.75]. Sketch of proof: Define $U=\{x:L_{x}(0)<\infty\}$ and the sequence of functions

Second claim. To find $f^{*}$ we use the Kuratowski-Ryll-Nardzewski Selection Theorem [AB94, Theorem 14.86] which states that if the (possibly multi-valued) map

is a weakly measurable correspondence (see [AB94, Chapter 14]) with nonempty closed values in a Polish space, then it admits a measurable selector. The only non-trivial fact is the weak measurability of $M$, which means that we need to check that the set

is measurable in $U$ whenever $V$ is open. Any open set $V$ admits a countable dense subset $\widehat{V}$, which means that

is measurable due to the continuity of $L_{x}$ in $\alpha$. We conclude that

For the reader’s convenience, we link Wasserstein gradient flows to classical gradient flows.

Let $\Theta=(\theta_{1},\dots,\theta_{m})$ and $f_{\Theta}(x)=\frac{1}{m}\sum_{i=1}^{m}\phi(\theta_{i};x)$. Then

Thus if the parameters $\theta_{i}$ evolve by the law $\dot{\theta}_{i}=-m\,\nabla_{\theta_{i}}\mathcal{R}(\Theta)$ for all $i=1,\dots,m$, then their distribution $\pi^{m}:=\frac{1}{m}\sum_{i=1}^{m}\delta_{\theta_{i}}$ satisfies the transport equation

by the flow map representation. This is precisely the PDE formulation of the 2-Wasserstein gradient flow. ∎

Now we show that gradient flow of $\mathcal{R}$ exists for any initial condition $\pi_{0}\in\mathcal{P}_{2}(\overline{\Theta})$ and that the omni-directionality of the initial measure is preserved along the gradient flow evolution (for finite time). Except for technical issues stemming from the lack of regularity in ReLU activation, the analysis follows [CB18b, Appendix B].

is a smooth cut-off function. In particular, note that

is Lipschitz continuous with Lipschitz constant $\|\rho\|_{L^{\infty}}$.

The first term on the right is bounded by

since $(a,w,b)\mapsto a\,\eta(a,w,b)$ is a positiively one-homogeneous function which is Lipschitz-continuous on the sphere and thus Lipschitz continuous. The second term on the right is

Tangency condition. We can write $\Theta=\{(a,w,b)\>|\>-a^{2}+|w|^{2}+b^{2}>0\}$. Then the normal to $\partial\Theta$ is parallel to $\nabla(-a^{2}+|w|^{2}+b^{2})=2(-a,w,b)$. Fix $(a,w,b)$ and consider any $x$ such that $w^{T}x+b\neq 0$. Since $\eta\equiv 1$ close to $\partial\Theta$, we have

since $\sigma$ is positively one-homogeneous. The equality also holds trivially at $x$ for which $w^{T}x+b=0$, so in particular after integration. ∎

In the calculus of variations (which encompasses the study of gradient flows), different notions of convexity play a key role. In vector spaces, convexity is a condition along straight lines, which (at least for Hilbert spaces) are the same as length-minimizing curves. The natural generalization to (geodesically complete) metric spaces is to consider the notion of convexity where a functional is ‘convex’ if it is convex along constant-speed length minimizing geodesics. The analogy is particularly strong in $2$-Wasserstein space, which carries the formal structure of a Hilbert manifold, see [Ott01] or [Vil08, Chapter 15]. This concept of convexity is referred to as displacement convexity and has been recognized since [McC97] as a useful notion when considering gradient flows.

The functionals we consider are not convex in Wasserstein space – in fact, convergence to a global minimizer is not guaranteed. For the existence of gradient flows, a weaker concept suffices. Recall that the functional $\mathcal{R}$ is called $\lambda$-displacement convex, if the following holds: If $\pi_{t}$ is a geodesic in Wasserstein space, then

By analogy with the smooth Euclidean case, we can think of the condition as a lower bound on the Hessian $D^{2}\mathcal{R}\geq-\lambda I$. Convexity corresponds to -convexity and uniform convexity to $\lambda$-convexity for $\lambda<0$.

It suffices to show that $h(t):=\mathcal{R}(\pi_{t})$ is $\lambda\,W_{2}^{2}(\pi_{0},\pi_{1})$-convex on $$, i.e.$h^{\prime\prime}\geq-\lambda\,W^{2}_{2}(\pi_{0},\pi_{1})$. We prove a stronger statement, namely that$h^{\prime}$is$\lambda\,W_{2}^{2}(\pi_{0},\pi_{1})$-Lipschitz, which can be thought of as a type of Hessian bound from both sides instead of just one side. Note that

Step 2. Existence of the gradient flow follows directly from [AGS08, Theorem 11.2.1].

Step 3. In this step, we show that the gradient flow of $\mathcal{R}$ preserves the cone $\Theta$, i.e. if $\pi_{0}\in\mathcal{P}_{2}(\overline{\Theta})$, then $\pi_{t}\in\mathcal{P}_{2}(\overline{\Theta})$ for all $t\geq 0$. Note that the flow field

is Lipschitz-continuous in $\theta$ with a uniform Lipschitz constant for all times by Lemma A.2. Like in Section 2.4, we find that $\pi_{t}=X(t,\cdot)_{\sharp}\pi_{0}$ where $X$ is the flow map defined by

It thus suffices to show that $X(t,\theta)\in\Theta$ for every $\theta\in\Theta$. This is immediate since

is parallel to $\partial\Theta$ on the boundary since $\nabla_{\theta}\phi(\theta,x)$ is parallel to $\partial\Theta$ at $\theta$ whenever it is defined – see also Section 2.5. ∎

We now prove that the gradient flow preserves the omni-directionality of measures (for finite positive time).

Since we can solve the ordinary differential equation backwards in time for any $\theta\in\Theta$, $X(t,\cdot):\Theta\to\Theta$ is a bi-Lipschitz homeomorphism.

Finally, we note that $X(t,\lambda\theta)=\lambda X(t,\theta)$ for all $\lambda,t>0$ since

due to the homogeneity of $\phi$. Thus, the flow preserves half-rays and cones.

Proof of omni-directionality. Consider an open cone $C\subseteq\Theta$. Then by [Amb08, Lemma 4], we have

since also $X(t,\cdot)^{-1}(C)$ is an open cone in $\Theta$. ∎

The analysis of the flow map shows more. Since rays are preserved, the projected measures $P^{S^{d+1}}_{\sharp}(\pi_{t})$ on the unit sphere $S^{d+1}\cap\Theta$ evolve by the continuity equation

where $\nabla^{S^{d+1}}f(\theta)=(I-\theta\otimes\theta)\nabla f(\theta)$ is the tangential gradient to the unit sphere. In particular, if $P^{S^{d+1}}_{\sharp}\pi_{0}$ has a density $\rho_{0}$ with respect to the uniform measure on $S^{d+1}\cap\Theta$, then $P^{S^{d+1}}_{\sharp}\pi_{t}$ has a density $\rho_{t}$ and

Next we show that the second moment of $\pi_{t}$ grows at most sublinearly in time.

Note that $\pi_{t}$ has finite second moments, so the quadratic test function $|\cdot|^{2}$ in the variational formulation is admissible. If $\pi_{0}$ is compactly supported, then so is $\pi_{t}$ for all $t>0$ by the flow map representation and the linear growth of the flow field at infinity. In this situation, the identity is obvious. Otherwise, the argument is easily justified by using approximating test functions $\eta(R-|\theta|)\,|\theta|^{2}$ where $\eta$ is a smooth cutoff function satisfying $\eta^{\prime}\geq 0$, $\eta(z)=0$ for $z\leq 0$ and $z=1$ for $z\geq 1$. Thus for every $0<T<t$ we have

Since $\mathcal{R}(\pi_{t})$ is monotone decreasing and bounded from below (by zero), $\mathcal{R}(\pi_{t})$ converges to a limit. Thus, for every $\varepsilon>0$ we can choose $T>0$ such that $0<\mathcal{R}(\pi_{T})-\mathcal{R}(\pi_{t})<\varepsilon$ for every $t>T$ and hence

Note that the proof applies in great generality to models with a linear structure. The next proof concerns the exponential growth of second moments if the velocity potential converges to a non-trivial limit. It is adapted from [CB18b] and repeated in this context for the reader’s convenience. The following proof is the only point at which the Morse-Sard property is used in this article. It is also the only argument which hinges on the omni-directionality of the initial parameter distribution.

$g(\theta)<-\varepsilon\,|\theta|^{2}$ for all $\theta\in C$, and

$\langle\nabla g,\nu_{C}\rangle>0$ does not vanish on $\partial C$ where $\nu$ is the inner normal vector to $\partial C$.

Using Assumption (M-S), $C$ can for example be chosen as

for some $t\in(\varepsilon,2\varepsilon)$. We define the localized second moments

Since $\partial C\cap S^{d+1}$ is compact and $\nabla(\delta_{\pi}\mathcal{R})\to\nabla g$ locally uniformly, we find that there exists $T_{0}>0$ such that $\langle\nabla(\delta_{\pi}\mathcal{R})(\pi_{t};\cdot),\nu_{C}\rangle>0$ on $\partial C$ for all $t>T_{0}$. In particular, no mass flows out of $C$ after time $T_{0}$: If $X_{\theta}(T_{0})\in C$ then also $X_{\theta}(t)\in C$ for $t>T_{0}$. Thus

Secondly since $(\delta_{\pi}\mathcal{R})(\pi_{t},\cdot)\to g$ locally uniformly, there exists $T_{1}>0$ such that

for all $t>T_{1}$. Without loss of generality, we assume that $T_{1}=T_{0}$. In particular

for $t>T_{0}$ and $\theta\in X(T_{0})^{-1}(C)$, using the positive two-homogeneity of $(\delta_{\pi}\mathcal{R})$. Thus

for $t>T_{0}$ and $\theta\in X(T_{0})^{-1}(C)$, and consequently

If $\pi_{0}$ is omni-directional, then so is $\pi_{T_{0}}$ and $N_{C}(\pi_{T_{0}})>0$. ∎

Morally, assumption (M-S) is used to control the sign of a boundary flux term. Without an assumption of this type, the term would at most be asymptotically non-negative. It would be necessary to control the size of the boundary term by a volume contribution and its asymptotic behavior.

A.2. Proofs from Section 3

We use these results to prove the main theorem.

Step 1. Since $\delta_{\pi}\mathcal{R}$ and $\nabla\delta_{\pi}\mathcal{R}$ are positively two- and one-homogeneous respectively, we find that $(\delta_{\pi}\mathcal{R})(\pi,0)=0$ and $\nabla(\delta_{\pi}\mathcal{R})(\pi,0)=0$ for any $\pi\in\mathcal{P}_{2}$. According to Lemma A.2, we may assume a uniform Lipschitz bound on $\nabla(\delta_{\pi}\mathcal{R})(\pi_{t},\cdot)$. This implies a uniform Lipschitz bound also on $\delta_{\pi}\mathcal{R}$ on bounded sets. We thus conclude that $(\delta_{\pi}\mathcal{R})(\pi_{t},\cdot)$ and $\nabla(\delta_{\pi}\mathcal{R})(\pi_{t},\cdot)$ have convergent subsequences, since Lipschitz-space embeds compactly into the space of continuous functions by the Arzelà-Ascoli theorem.

Step 2. First, assume that $\Omega^{lim}$ consists of only one element $(g,V)$. Then either $g\not\equiv 0$ or $g\equiv 0$. In the first case, we have by Lemma 3.4 that $N(\pi_{t})$ grows exponentially in time, contradicting Lemma 3.3. In the second case, we use homogeneity to show that

because $\widetilde{\mathcal{R}}(f^{*})=\inf_{\pi}\mathcal{R}(\pi)$. Since $N$ is bounded from below by zero, it cannot decrease linearly at a fixed non-zero rate for all large arguments. We conclude that

Step 3. Now, assume that $\lim_{t\to\infty}\mathcal{R}(\pi_{t})=\widetilde{\mathcal{R}}(f^{*})$, i.e.

where $L$ is the augmented loss function discussed in Section 2.2. Since the first integrand is always larger than the second one, their difference is positive and thus we conclude that

Thus $\lim_{n\to\infty}\delta_{\pi}\mathcal{R}(\pi_{t_{n}},\cdot)\equiv 0$. By compactness, we know that $\Omega^{lim}$ is non-empty, and we have showed that for any sequence, we can extract a subsequence for which $g=0$. Since locally uniform convergence is generated by a topology, this means that $g$ is the only possible limit point. The same argument can be used for the gradient. ∎

Note that omni-directionality is only used to exclude the case that $g\not\equiv 0$, whereas $g\equiv 0$ is only an admissible limit if $\mathcal{R}(\pi_{t})$ decays to MBR. This corresponds to the fact that $\delta_{\pi}\mathcal{R}(\pi,\cdot)\equiv 0$ if and only if $\pi$ is a global minimizer of $\mathcal{R}$, see [CB18b, Proposition 3.1].

Omni-directionality of the initial condition (IC2) and the Morse-Sard property are both involved only in excluding a unique limit $g\not\equiv 0$. They could therefore be replaced for example by the zero-limit assumption

If $(\delta_{\pi}\mathcal{R})(\pi_{t},\cdot)\to g$ locally uniformly on $\overline{\Theta}$ and $g\not\equiv 0$, then

The following two proofs establish the corollaries to the main theorem concerning minimizers.

If we assume in addition that there exists a probability measure $\pi_{\infty}$ and a sequence of times $t_{n}\to\infty$ such that $\lim_{n\to\infty}W_{2}(\pi_{t_{n}},\pi_{\infty})=0$, then we find by [Vil08, Theorem 6.9] that

Thus $\pi_{\infty}$ minimizes $\mathcal{R}$. ∎

Consider the measures $\mu_{t}$ on $S^{d+1}\cap\Theta$ defined by

for $g\in C(\overline{S^{d+1}\cap\Theta})$, or equivalently

where $P^{S^{d+1}}:\Theta\to S^{d+1}\cap\Theta$ is the canonical projection. While $P^{S^{d+1}}$ is undefined at $\theta=0$, the projection of the measure with weight $|\theta|^{2}$ is well-defined. Under the moment bound assumption, the measures $\mu_{t}$ are uniformly bounded and

$\mu_{\infty}=0$. In this case $f^{*}\equiv 0$ and $\pi_{\infty}:=\delta_{\theta=0}$ is a risk minimizer.

$\mu_{\infty}\neq 0$. In this case, we define the dilation map $D=\sqrt{\mu_{\infty}(\Theta)}\,I$ and

for all $x\in U$, so $\pi_{\infty}$ is a risk minimizer.

Thus the projection to the unit sphere adds compactness beyond the moment bound.

Appendix B The Morse-Sard Property

The result is due to Morse for $m=1$ [Mor39] and Sard for general $m\geq 1$ [Sar42]. It is easy to extend the result to sufficiently differentiable manifolds, and there are more precise statements available for the Hausdorff measure of $f(S_{\nu})$ where

Morse-Sard theorems are known to fail in infinite dimension even for infinitely smooth maps, unless additional assumptions are imposed. Under weak conditions, however, the set of functions for which the Theorem holds is dense in the $C^{0}$-topology [EM68].

Due to its fundamental importance, some effort has been made to establish Morse-Sard type properties in other function classes. Among these are

Morse-Sard theorems in classes of weakly differentiable functions [Fig08, BKK13, BKK15, dP01, KK18]. Here the relation between differentiability and integrability may even be chosen low enough to ensure continuity, but not classical differentiability of the functions under consideration.

Morse-Sard theorems for the distance from a submanifold [Rif04] or more generally Lipschitz functions which are given as suprema of smooth functions over suitable index sets [BDDR16].

Morse-Sard theorems for subanalytic functions, see [BDL06].

Morse-Sard theorems for dc functions in two dimensions [PZ06]. A function is dc if it can be written as differences of two convex functions. In particular, every $C^{2}$-function with bounded second derivatives is dc. Thus the result cannot be generalized to higher dimensions.

In non-smooth function classes, a notion of gradient almost everywhere (with respect to a suitable Hausdorff measure) or a sub-differential is used.

We show below that (M-S) holds unconditionally in dimension $d=2$. In [CB20], the Morse-Sard theorem for subanalytic functions from [BDL06] has successfully been used to establish a condition of Morse-Sard property in a very similar application. The subanalytic function that the authors consider in [CB20] is a finite sum of ReLU-like terms and the subanalyticity stems from the finiteness of the sum. The number of summands corresponds to the number of data samples in an empirical measure. The approach is therefore incompatible with assumption (P4) for the ReLU case.

While it does not apply in our situation, we briefly sketch the result and its application in a similar situation. Consider $f(x)=\sum_{i=1}^{m}a_{i}(w_{i}^{T}x+b_{i})_{+}$. Then $f$ is only Lipschitz-smooth and not $C^{1}$, but has the Morse-Sard property due to its sub-analyticity.

We recall the following Morse-Sard theorem.

We show that this applies to finite ReLU networks.

$f(x)=\sum_{i=1}^{m}a_{i}(w_{i}^{T}x+b_{i})_{+}$ is continuous and sub-analytic.

Continuity is clear. Let $(x,y)\in G_{f}$. Since $f$ is affine linear on the set $\{x\>|\>w_{i}^{T}x+b_{i}\neq 0\>\forall\ i\}$, the graph is a plane and thus sub-analytic here. Now assume that $w_{1}^{T}x+b_{1}=0$ and $w_{i}^{T}x+b_{i}\neq 0$ for all $i>1$. Then, locally after a rotation and translation we have

The graph of this function is sub-analytic since

The case when more terms $w_{i}^{T}x+b_{i}$ vanish can be treated similarly, but is somewhat tedious to write out. It is, however, crucial that the sum is finite to ensure that there are at most finitely many sets to be united and intersected. ∎

B.2. ReLU Geometry on the Sphere

We can exploit the special geometry of the problem to reduce the dimension slightly. We can write

the only critical value is (possibly) zero. On the other hand, consider $(a,w)$ in $\partial\Theta\cap S^{d+1}$. We note that

B.3. A mild counterexample

We show that functions with similar structural properties as $g$ (or $h$) may not satisfy a Morse-Sard property in dimension $d\geq 8$.

The space of Barron functions is discussed in detail in [EMW18, EMW19b, EW20a] and [EW20b, Appendix A]. The space is named after Andrew Barron who first established that a large class of functions could be represented in such a way. We cite a simplified version of Barron’s main theorem.

Using the identity $\widehat{\partial_{j}f}=i\,\xi_{j}\widehat{f}$ and Parseval’s identity, we compute

In dimension $d\geq 7$ there exist Barron functions which do not have the Morse-Sard property.

References