A Dynamical Central Limit Theorem for Shallow Neural Networks

Theoretical analyses of neural networks aim to understand their computational and statistical advantages seen in practice. On the computation side, the training of neural networks often succeed despite being a non-convex optimization problem known to be hard in certain settings . On the statistics side, neural networks often generalize well despite having large numbers of parameters . In this context, the notion of over-parametrization has been useful, by providing insights into the optimization and generalization properties as the network widths tend to infinity . In particular, under appropriate scaling, one can view shallow (a.k.a. single-hidden-layer or two-layer) networks as interacting particle systems that admit a mean-field limit. Their training dynamics can then be studied as Wasserstein Gradient Flows , leading to global convergence guarantees in the mean-field limit under certain assumptions. On the statistics side, such an approach lead to powerful generalization guarantees for learning high-dimensional functions with hidden low-dimensional structures, as compared to learning in Reproducing Kernel Hilbert Spaces (RKHS) . However, since ultimately we are concerned with neural networks of finite width, it is key to study the deviation of finite-width networks from their infinite-width limits, and how it scales with the width $m$. At the random initial state, neurons do not interact and therefore a standard Monte-Carlo (MC) argument shows that the fluctuations in the underlying measure scale as $m^{-1/2}$, which we refer to as the Central Limit Theorem (CLT) scaling. As optimization introduces complex dependencies among the parameters, the key question is to understand how the fluctuation evolves during training. To make this investigation tractable, we aim to obtain insight on an asymptotic scale as the width grows, and focus on the evolution in time. An application of Grönwall’s inequality shows that this asymptotic deviation remains bounded at all finite time , but the dependence on time is exponential, making it difficult to assess the long-time behavior.

The main focus of this paper is to investigate this question in-depth, by analyzing the interplay between the deviations from the mean-field limit and the gradient flow dynamics. First, we prove a dynamical CLT to capture how the fluctuations away from the mean-field limit evolve as a function of training time to show that the fluctuations remain on the initial $m^{-1/2}$-scale for all finite times. Next, we examine the long-time behavior of the fluctuations, proving that, in several scenarios, the long-time fluctuations are controlled by the error of Monte-Carlo resampling from the limiting measure. We focus on two main setups relevant for supervised learning and scientific computing: the unregularized case with global convergence of mean-field gradient flows to minimizers that interpolate the data, and the regularized case where the limiting measure has atomic support and is nondegenerate. In the former setup, we prove particularly that the fluctuations eventually vanish in the CLT scaling. These asymptotic predictions are complemented by empirical results in a teacher-student model.

This paper continues the line of work initiated in that studies optimization of over-parameterized shallow neural networks under the mean-field scaling. Global convergence for the unregularized setting is discussed in . In the regularized setting, establishes global convergence in the mean-field limit under specific homogeneity conditions on the neuron activation. Other works that study asymptotic properties of wide neural networks include , notably investigating the transition between the so-called lazy and active regimes , corresponding respectively to linear versus nonlinear learning. Our focus is on the dynamics under the mean-field scaling, which encompasses the active, nonlinear regime.

A relevant work concerning the sparse optimization of measures is , where under a different metric for gradient flow and additional assumptions on the nature of the minimizer, it can be established that fluctuations vanish for sufficiently large $m$. Our results are only asymptotic in $m$ but apply to broader settings in the context of shallow neural networks. Concerning the next-order deviations of finite neural networks from their mean-field limit, show that the scale of fluctuations is below that of MC resampling for unregularized problems using non-rigorous arguments. provides a CLT for the fluctuations at finite time under stochastic gradient descent (SGD) and proves that the fluctuations decay in time in the case where there is a single critical point in the parameter space. Our focus is on the long-time behavior of the fluctuations in more general settings. Another relevant topic is the propagation of chaos in McKean-Vlasov systems, which study the deviations of randomly-forced interacting particle systems from their infinite-particle limits . In particular, a line of work provides uniform-in-time bounds to the fluctuations in various settings , but the conditions are not applicable to shallow neural networks. Concurrently to our work, studies quantitative propagation of chaos of shallow neural networks trained by SGD, but the bound grows exponentially in time, and therefore cannot address the long-time behavior of the fluctuations.

Learning with neural networks exhibits the phenomenon that generalization error can decrease with the level of overparameterization . proposes a bias-variance decomposition that contains a variance term initialization in optimization. They show in experiments that this term decreases as the width of the network increases, and justifies this theoretically under the strong assumption that model parameters remain Gaussian-distributed in the components that are irrelevant for the task, which does not hold in the scenario we consider, for example. provides scaling arguments for the dependence of this term on the width of the network. Our work provides a more rigorous analysis of the dependence of this term on the width of the network and training time.

Background

As many of our results hold for general models of the form (1), we will invoke Assumption 2.1 only when needed. We shall also assume the following:

$\Omega$ is compact; $D$ is a Euclidean space (or a subset thereof); $\varphi(\bm{\theta},\bm{x})$ is twice differentiable in $\bm{\theta}$; $\nabla_{\bm{\theta}}\nabla_{\bm{\theta}}\varphi(\bm{\theta},\bm{x})$ is Lipschitz in $\bm{\theta}$, uniformly in $\bm{x}$.

As observed in , a model of the form (1) can be expressed in integral form in terms of a probability measure over $D$ as $f^{(m)}=f[\mu^{(m)}]$, where we define

and $\mu^{(m)}$ is the empirical measure of the parameters $\{\bm{\theta}_{i}\}_{i=1}^{m}$:

Suppose we are given a dataset $\{(\bm{x}_{l},y_{l})\}_{l=1}^{n}$, which can be represented by an empirical data measure $\hat{\nu}=\tfrac{1}{n}\sum_{l=1}^{n}\delta_{\bm{x}_{l}}$, and $y_{l}=f_{*}(\bm{x}_{l})$ are generated by an target function $f_{*}$ that we wish to estimate using least-squares regression. A canonical approach to this regression task is to consider an Empirical Risk Minimization (ERM) problem of the form

where $\mathcal{P}(D)$ is the space of probability measures on $D$, $\|f-f_{*}\|_{\hat{\nu}}^{2}:=\int_{\Omega}|f(\bm{x})-f_{*}(\bm{x})|^{2}\hat{\nu}(d\bm{x})$ denotes the $L_{2}$ function reconstruction error averaged over the data, and $\lambda\int_{D}r(\bm{\theta})\mu(d\bm{\theta})$ is some optional regularization term. While we can allow $r$ to be a general convex function, in Section 3.3 we will motivate a choice of $r$ in the shallow neural networks setting that is related to the variation norm or Barron norm of functions.

2 Approximation and Optimization with a Finite Number of Neurons

Integral representations with a probability measure such as those defined in (2) are amenable to efficient approximation in high dimensions via Monte-Carlo sampling. Namely, if the parameters $\bm{\theta}_{i}$ in $f^{(m)}$ are drawn i.i.d. from an underlying measure $\mu$ on $D$, then by the Law of Large Numbers (LLN), the resulting empirical measure $\mu^{(m)}$ converges $\mu$ almost surely, and moreover,

Such a Monte-Carlo estimator showcases the benefit of normalized integral representations for high-dimensional approximation, as the ambient dimension appears in the rate of approximation only through the term $\int_{D}\|\varphi(\bm{\theta},\cdot)\|_{\hat{\nu}}^{2}\mu(d\bm{\theta})$. In the case of shallow neural networks, this is connected to the variation norm or Barron norm of the function we wish to approximate (see Section 3.3 for details).

While the Monte-Carlo sampling strategy above can be seen as a ‘static’ approximation of a function representable as (2), it also gives rise to an efficient algorithm to optimize (4). Indeed, in terms of the empirical distribution $\mu^{(m)}$, the loss $\mathcal{L}(\mu^{(m)})$ becomes a function of the parameters $\{\bm{\theta}_{i}\}_{i=1}^{m}$, which we can seek to minimize by adjusting the parameters:

In the shallow neural network setting, with suitable choices of the function $r$, the regularization term corresponds to weight decay over the parameters.

3 From Particle to Wasserstein Gradient Flows

where we have defined $C_{f}=\frac{1}{2}\|f\|_{\hat{\nu}}^{2}$, and

Performing GD on ${L}$ amounts to discretizing in time the following ODE system that governs the evolution of for $\{\bm{\theta}_{i}\}_{i=1}^{m}$:

Heuristically, the ‘particles’ $\bm{\theta}_{i}$ perform GD according to the potential $V(\bm{\theta},\mu_{t}^{(m)})$ which itself evolves, depending on the particles positions through their empirical measure. Such dynamics can also be expressed in terms of the empirical measure via the continuity equation:

4 Law of Large Numbers and Mean-Field Gradient Flow

From now on, we assume that the particle gradient flow is initialized in the following way:

The ODE (9) is solved for the initial condition $\bm{\theta}_{i}(0)=\bm{\theta}_{i}^{0}$, with $\bm{\theta}_{i}^{0}$ drawn i.i.d. from a compactly supported measure $\mu_{0}\in\mathcal{P}(D)$ for each $i=1,\ldots,m$. Hence, $\mu^{(m)}_{0}(d\bm{\theta})=\tfrac{1}{m}\sum_{i=1}^{m}\delta_{\bm{\theta}_{i}^{0}}(d\bm{\theta})$.

The solution to this equation can be understood via the representation formula

Using expression (10) for $V$ as well as (13), this equation can be written in closed form explicitly as

It is easy to see that this equation is itself a gradient flow since it is the continuous-time limit of a proximal scheme (mirror descent), which we state as:

5 Long-Time Properties of the Mean-Field Gradient Flow

In the shallow neural networks setting, a series of earlier works has established that under certain assumptions $\mu_{t}$ will converge to a global minimizer of the loss functional $\mathcal{L}$. In particular, studies global convergence for the regularized loss $\mathcal{L}$ under homogeneity assumptions on $\hat{\varphi}$, and considers modified dynamics using double-lifting. Here, to study the long time behavior of the fluctuations, we will often work with the following weaker assumptions:

The solution to (15) exists for all time, and has a limit:

The limiting $\bm{\Theta}_{\infty}$ is a local minimizer of (18).

and $\mu_{t}\rightharpoonup\mu_{\infty}$ weakly as $t\to\infty$, with $\mu_{\infty}$ satisfying

We prove this proposition in Appendix B. Here, $\nabla\nabla V(\bm{\Theta}_{\infty}(\bm{\theta}),\mu_{\infty})$ denotes

which will become useful in Section 3.2 when we analyze the long time properties of the fluctuations around the mean-field limit.

Assumptions 2.5 and 2.6 impose conditions on the initial measure $\mu_{0}$ . While the convergence of gradient flows in finite-dimensional Euclidean space to local minimizers is guaranteed under mild assumptions , its infinite-dimensional counterpart, Assumption 2.6, may require further technical assumptions, left for future study. Also, while Assumption 2.5 implies that $\mu_{\infty}$ is a stationary point of (12), Assumption 2.6 does not imply that $\mu_{\infty}$ minimizes $\mathcal{L}$.

Fluctuations from Mean-Field Gradient Flow

The main goal of this section is to characterize the deviations of finite-particle shallow networks from their mean-field evolution, by first deriving an estimate for $f^{(m)}_{t}-f_{t}$ for $t\geq 0$ (Section 3.1), and then analyzing its long-time properties (Section 3.2). In Section 3.3, we then motivate a choice of the regularization term in (4) that controls the bound on the long-time fluctuations derived in Section 3.2, and which is also connected to generalization via the variation norm or Barron norm of functions.

By the static Central Limit Theorem (CLT) we know that, if we draw the initial values of the parameters $\bm{\theta}_{i}$ independently from $\mu_{0}$ as specified in Assumption 2.3, $g^{(m)}_{0}$ has a limit as $m\to\infty$, leading to estimates similar to (5) with $\mu^{(m)}$ and $\mu$ replaced by the initial $\mu_{0}^{(m)}$ and $\mu_{0}$, respectively. For $t>0$, however, this estimate is not preserved by the gradient flow: the static CLT no longer applies and needs to be replaced by a dynamical variant . Next, we derive this dynamical CLT in the context of neural network optimization.

To this end let us define the discrepancy measure $\omega^{(m)}_{t}$ such that

Hence, we will first establish how the limit of $\omega^{(m)}_{t}$ as $m\to\infty$ evolves over time. This can be done by noting that the representation formula (13) implies that

where $\bm{\Theta}^{(m)}_{t}$ solves (15) with $\mu_{0}$ replaced by $\mu_{0}^{(m)}$. Defining

As shown in Appendix C.1, we can take the limit $m\to\infty$ of this formula to obtain:

Here $\omega_{0}$ is the Gaussian measure with mean zero and covariance

with initial condition $\bm{T}_{0}=0$ and where $\bm{\Theta}_{t}$ solves (14) and $\nabla\nabla V(\bm{\Theta}_{t}(\bm{\theta}),\mu_{t})$ is a shorthand for

A direct consequence of this proposition and formula (27) is:

It is interesting to comment on the origin of both terms at the right hand side of (31) and, consequently, (35). The first term captures the deviations induced by fluctuations of $\mu_{0}^{(m)}$ around $\mu_{0}$ assuming that the flow $\bm{\Theta}^{(m)}_{t}$ is unaffected by these fluctuations, and remains equal to $\bm{\Theta}_{t}$. In particular, this term is the one we would obtain if we were to resample $\mu_{t}^{(m)}$ from $\mu_{t}$ at every $t\geq 0$, i.e. use $\bar{\mu}_{t}^{(m)}=m^{-1}\sum_{i=1}^{m}\delta_{\bar{\bm{\theta}}_{t}^{i}}$ with $\{\bar{\bm{\theta}}_{t}^{i}\}_{i=1}^{m}$ sampled i.i.d. from $\mu_{t}$, so that $\bm{\Theta}^{(m)}_{t}$ is identical to $\bm{\Theta}_{t}$ in (28). In this case, the limiting discrepancy measure $\bar{\omega}_{t}$ would simply be given by

while the associated deviation in the represented function would read

The second term at right hand side of (31) and (35) captures the deviations to the flow $\bm{\Theta}_{t}$ in (15) induced by the perturbation of $\mu_{0}$, i.e. how much $\bm{\Theta}^{(m)}_{t}$ differs from $\bm{\Theta}_{t}$ in (28). In the limit as $m\to\infty$, these deviations are captured by the solution $\bm{T}_{t}$ to (33), as is apparent from (30).

The difference between $g_{t}$ and $\bar{g}_{t}$ can also be quantified via the following Volterra equation, which can be derived from Proposition 3.1 and relates the evolution of $g_{t}$ to that of $\bar{g}_{t}$.

Here $\bar{g}_{t}$ is given in (37) and we defined

This corollary is proven in Appendix C.2. In a nutshell, (38) can be established using Duhamel’s principle on (33) by considering all terms at the right hand side except the first as the source term (hence the role of $J_{t,s}$) and inserting the result in (35).

2 Long-Time Behavior of the Fluctuations

Next, we study the long-time behavior of $g_{t}$ and, in particular, evaluate

It is useful to start by considering an idealized case, namely when the initial conditions are sampled as in Assumption 2.3 with $\mu_{0}=\mu_{\infty}$. In that case, there is no evolution at mean field level, i.e. $\bm{\Theta}_{t}(\bm{\theta})=\bm{\Theta}_{\infty}(\bm{\theta})=\bm{\theta}$, $\mu_{t}=\mu_{\infty}$, and $f_{t}=f_{\infty}=\int_{D}\varphi_{\infty}(\bm{\theta},\cdot)\mu_{\infty}(d\bm{\theta})$, but the CLT fluctuations still evolve. In particular, it is easy to see that the Volterra equation in (38) for $g_{t}$ becomes

Here $\Gamma^{\infty}_{t-s}(\bm{x},\bm{x}^{\prime})$ is the Volterra kernel obtained by solving (40) with $\nabla\nabla V(\bm{\Theta}_{t}(\bm{\theta}),\mu_{t})$ replaced by $\nabla\nabla V(\bm{\theta},\mu_{\infty})$ and inserting the result in (39) with $\bm{\Theta}_{t}(\bm{\theta})=\bm{\theta}$ and $\mu_{0}=\mu_{\infty}$,

and $\bar{g}_{\infty}$ is the Gaussian field with variance

From (23) in Proposition 2.7 we know that $\nabla\nabla V(\bm{\theta},\mu_{\infty})$ is positive semidefinite for $\mu_{\infty}$-almost all $\bm{\theta}$. As a result, we prove in D.1 that the Volterra kernel (43) viewed as an operator on functions defined on $\Omega\times[0,T]$ is positive semidefinite. Therefore, we have

Under Assumptions 2.2, 2.3, 2.5 and 2.6, with $\mu_{0}=\mu_{\infty}$ and $\mu_{\infty}$ as specified in Proposition 2.7, we have

This theorem indicates that, if we knew $\mu_{\infty}$ and could sample initial conditions for the parameters from it, it would still be favorable to train these parameters as this would reduce the approximation error. Of course, in practice we have no a priori access to $\mu_{\infty}$, and so the relevant question is whether (46) also holds if we sample initial conditions from any $\mu_{0}$ such that Proposition 2.7 holds.

In light of (35), one way to address this question is to study the long-time behavior of $\bm{T}_{t}$. In the setup without regularization ($\lambda=0$), we can do so by leveraging existing results that, under certain assumptions, the mean-field gradient flow converges to a global minimizer which interpolates the training data points exactly . In this case, the following theorem shows that we can actually obtain stronger controls on the fluctuations than (46), which we prove in Appendix D.2.

Consider the ERM setting with $\lambda=0$ and under Assumptions 2.2, 2.3 and 2.5. Suppose that as $t\to\infty$, $\mu_{t}$ converges to a global minimizer $\mu_{\infty}$ that interpolates the data, i.e. the function $f_{\infty}=\int_{D}\varphi(\bm{\theta},\cdot)\mu_{\infty}(d\theta)$ satisfies

and, furthermore, the convergence satisfies

if Assumption 2.1 also holds, i.e., in the shallow neural network setting, we further have

if $\mu_{0}=\mu_{\infty}$, then $\|g_{t}\|_{\hat{\nu}}$ decreases monotonically in $t$.

Hence, in the shallow neural networks setting and under these assumptions, the fluctuations will eventually vanish in the $O(m^{-1/2})$ scale of CLT. Note that for (48) to hold, it is sufficient that $\mathcal{L}(\mu_{t})$ decays at an asymptotic rate of $O(t^{-\alpha})$ with $\alpha>4$. For instance, proves that in an ERM setting where the size of the training dataset is no larger than the input dimension (i.e. $n\leq d$), the loss converges to zero at a linear rate, which will satisfy the condition (48). We leave the search for weaker sufficient conditions for future work.

When the limiting measure $\mu_{\infty}$ does not necessarily interpolate the training data, such as when regularization is added, we can proceed with the analysis of the long-time behavior of $\bm{T}_{t}$ under the following assumption on the long-time behavior of the curvature:

Let $\Lambda_{t}(\bm{\theta})$ denote the smallest eigenvalue of the tensor $\nabla\nabla V(\bm{\Theta}_{t}(\bm{\theta}),\mu_{t})$ defined in (34) and assume that for a constant $C$ (to be specified in Appendix D.3) such that

This theorem is proven in Appendix D.3. To intuitively understand (50), note that we know from (23) in Proposition 2.7 that $\Lambda_{t}(\bm{\theta})\to 0$ $\mu_{0}$-almost surely as $t\to\infty$. Condition (50) can therefore be satisfied by having $\Lambda_{t}(\bm{\theta})$ converge to zero sufficiently fast in the regions of $D$ where it is negative, or having the measure of these regions with respect to $\mu_{0}$ converge to zero sufficiently fast, or both.

Alternatively, in the regularized ($\lambda>0$) ERM setting, we can obtain the following result when the support of $\mu_{\infty}$ is atomic, as expected on general grounds :

Consider the ERM setting under Assumptions 2.2, 2.3 and 2.5. Suppose further that as $t\to\infty$, $\mu_{t}$ converges to $\mu_{\infty}$ satisfying

Then (46) holds with the “$\lim$" replaced by “$\limsup$" on its LHS.

Theorem 3.7 is proven in Appendix D.4 by analyzing directly the Volterra equation (38) and establishing that its solution coincides with that of (42) in the limit as $t\to\infty$, a property that we also expect to hold more generally than under the assumptions of Theorem 3.7. In fact, we prove in Appendix D.4 that (52) can be replaced by a weaker condition, (238). We also discuss the relation between Theorem 3.7 and the work of in Appendix D.4.3.

3 The Monte-Carlo Bound and Regularization

The bound (46) on the long-time fluctuations motivates us to control the term $\int_{D}\|\varphi(\bm{\theta},\cdot)\|_{\hat{\nu}}^{2}\mu_{\infty}(d\bm{\theta})$ using a suitable choice of regularization in (4). In the following, we restrict our attention to the shallow neural networks setting, and further assume that

where $\hat{K}_{M}=\max_{\bm{z}\in\hat{D}}\|\hat{\varphi}(\bm{z},\cdot)\|_{\hat{\nu}}^{2}$. Thus, we consider regularization with $r(\bm{\theta})=\frac{1}{2}c^{2}$, in which case (4) becomes

Interestingly, this choice of regularization leads to learning in the function space $\mathcal{F}_{1}$ (or alternatively, the Barron space ) associated with $\hat{\varphi}$, which is equipped with the variation norm (or the Barron norm) defined as

The space $\mathcal{F}_{1}$ contains any RKHS whose kernel is generated as an expectation over features $k(\bm{x},\bm{x}^{\prime})=\int_{\hat{D}}\hat{\varphi}(\bm{z},\bm{x})\hat{\varphi}(\bm{z},\bm{x}^{\prime})\hat{\mu}_{0}(d\bm{z})$ with a base measure $\hat{\mu}_{0}\in\mathcal{P}(\hat{D})$, but it provides crucial approximation advantages over such RKHS at approximating certain non-smooth, high-dimensional functions with hidden low-dimensional structure, giving rise to powerful generalization guarantees . This also motivates the study of overparametrized shallow networks with the scaling as in (1), as opposed to the NTK scaling of $m^{-1/2}$ .

To learn in $\mathcal{F}_{1}$, a canonical approach is to consider the ERM problem

By (55), this is indeed equivalent to (54). In Appendix E, we prove the following proposition, which shows that the measure obtained from (54) indeed has its $2$-norm controlled:

Under Assumptions 2.1, 2.2, and 3.8, $\mathcal{L}$ has no local minima and its global minimum value can only be attained at measures $\mu_{\lambda}\in\mathcal{P}(D)$ such that both $f_{\lambda}=\int_{D}\varphi(\bm{\theta},\cdot)\mu_{\lambda}(d\bm{\theta})$ and $c_{\lambda}=\int_{D}|c|\mu_{\lambda}(d\bm{\theta})=\left(\int_{D}|c|^{2}\mu_{\lambda}(d\bm{\theta})\right)^{1/2}\leq\gamma_{1}(f_{*})$ are unique, and

where $\hat{K}_{M}=\max_{\bm{z}\in\hat{D}}\|\hat{\varphi}(\bm{z},\cdot)\|_{\hat{\nu}}^{2}$.

Numerical Experiments

We first perform numerical experiments in a student-teacher setting, using a shallow teacher network as the target function to be learned by shallow student networks with different widths $m$ of the hidden layer. Both $\hat{D}$ and $\Omega$ are taken to be the unit sphere of $d=16$ dimensions, and we take $\hat{\varphi}(\bm{z},\bm{x})=\max(0,\langle\bm{z},\bm{x}\rangle)$. The teacher network has two neurons, $(c_{1},\bm{z}_{1})$ and $(c_{2},\bm{z}_{2})$, in the hidden layer, with $c_{1}=c_{2}=1$ and $\bm{z}_{1}$ and $\bm{z}_{2}$ sampled i.i.d. from the uniform distribution on $\hat{D}$ and then fixed across the experiments. We vary the width of the student network in the range of $m=128,256,512,1024$ and $2048$, with their initial $\bm{z}_{i}$’s sampled i.i.d. from the uniform distribution on $\hat{D}$. We consider two ways for initializing the $c_{i}$’s of the student networks: 1) Gaussian-initialization, where the $c_{i}$’s are sampled i.i.d. from $\mathcal{N}(0,1)$; and 2) zero-initialization, where each $c_{i}$ is set to be .

We train the student networks in two ways: using the population loss or the empirical loss. For the former scenario, the data distribution $\nu$ is chosen to be uniform on $\Omega$, which allows an analytical formula for the loss as well as its gradient. The student networks are trained by gradient descent under $L_{2}$ loss. Moreover, we rescale both the squared loss and the gradient by $d$ in order to adjust to the $\frac{1}{d}$ factor resulting from spherical integrals, and set the learning rate (which is the step size for discretizing (9)) to be $1$. The models are trained for $20000$ epochs. For each choice of $m$, we run the experiment $\kappa=20$ times with different random initializations of the student network. The average fluctuation of the population loss is defined as $\frac{1}{\kappa}\sum_{k=1}^{\kappa}\|f^{(m)}_{k}-\bar{f}^{(m)}\|_{\nu}^{2}$ for the population loss, with $\bar{f}^{(m)}=\frac{1}{\kappa}\sum_{k=1}^{\kappa}f^{(m)}_{k}$ being the averaged model, similar to the approach in . The other plotted quantities – loss, TV-norm and $2$-norm – are averaged across the $\kappa$ number of runs. The TV-norm (i.e., $1$-norm) and $2$-norm are defined as in Appendix 3.3.

The results for the scenario of training under the population loss are presented in Figures 1. As seen from Column 3 the average loss values remain similar over time for different choices of $m$, justifying the approximation by a mean-field dynamics. In the unregularized case with non-zero initialization, the fluctuation of the population loss (shown in Column 2) remains close to a $1/m$ scaling in roughly the first $10^{3}$ epochs, after which it decays faster for smaller $m$. Interestingly, this coincides with the tendency for the student neurons with $\bm{z}$ not aligned with the teacher neurons to slowly have their $|c|$ decrease to zero due to a finite-$m$ effect, which is also reflected in the decrease in TV-norm. Aside from this phenomenon, the fluctuations decay at similar rates for different choices of $m$, which is consistent with our theory, since their dynamics are governed by the same dynamical CLT. Also, when regularization is added, each student neuron becomes aligned with one of the teacher neurons in both $\bm{z}$ and $c$ after training; without regularization but using zero-initialization, after training, each student neuron either becomes aligned with one of the teacher neurons in $\bm{z}$ or has $c$ close to zero. Both of these choices result in lower TV-norms and $2$-norms compared to using non-zero initialization and without regularization.

Next, we consider the empirical loss scenario (ERM setting), using $n=32$ random vectors sampled i.i.d. from the uniform distribution $\nu$ on $\Omega$ as the training dataset, which then define the empirical data measure $\hat{\nu}(d\bm{x})=\frac{1}{n}\sum_{l=1}^{n}\delta_{\bm{x}_{l}}(d\bm{x})$. We use the full training dataset for computing the gradient at every iteration. The other training setups are the same as when the population loss is used. We additionally plot the average fluctuation of the training loss, defined as $\frac{1}{\kappa}\sum_{k=1}^{\kappa}\|f^{(m)}_{k}-\bar{f}^{(m)}\|_{\hat{\nu}}^{2}$.

The results for the scenario of training under the empirical loss is presented in Figures 2. Compared to the scenario of training under the population loss, we wee that in the unregularized cases, both the average training loss and the average fluctuation of the training loss decay to below $10^{-8}$ within $10^{3}$ iterations, and the latter observation is consistent with (49). In the regularized case, neither of them vanishes, but the average fluctuation of the training loss indeed remains below the asymptotic Monte-Carlo bound given in (46), whose analytical expression and numerical value in this setup (under the approximation of replacing $\mu_{\infty}$, $f_{\infty}$ and $\hat{\nu}$ by the target measure, the target function and $\nu$, respectively) are given in Appendix F. Regularization and zero-initialization have a weaker effect in aligning the student neurons with the teacher neurons after training compared to the scenario of training under the population loss, but they still result in lower TV-norm and $2$-norm, and moreover, lower average fluctuation and (slightly) lower average value of the population loss. This demonstrates their positive effects on both approximation and generalization.

2 Non-planted Case

The results are shown in Figure 3. We observe that the behaviors of the fluctuation are qualitatively similar to those found in Figure 1.

Conclusions

Here we studied the deviations of shallow neural networks from their infinite-width limit, and how these deviations evolve during training by gradient flow. In the ERM setting, we established that under different sets of conditions, the long-term deviation under the Central Limit Theorem (CLT) scaling is controlled by a Monte Carlo (MC) resampling error, giving width-asymptotic guarantees that do not depend on the data dimension explicitly. The MC resampling bound motivates a choice of regularization that is also connected to generalization via the variation-norm function spaces.

Our results thus seem to paint a favorable picture for high-dimensional learning, in which the optimization and generalization guarantees for the idealized mean-field limit could be transferred to their finite-width counterparts. However, we stress that these results are asymptotic, in that we take limits both in the width and time. In the face of negative results for the computational efficiency of training shallow networks , an important challenge is to leverage additional structure in the problem (such as the empirical data distribution , or the structure of the minimizers ) to provide nonasymptotic versions of our results, along the lines of or . Finally, another clear direction for future research is to extend our techniques to deep neural architectures, in light of recent works that consider deep or residual models .

Acknowledgements

This work benefited from discussions with Lenaic Chizat and Carles Domingo-Enrich, and the authors sincerely thank Jiaheng Chen for pointing out an error in Theorem 3.5 in the previous version of this manuscript. Z.C. acknowledges support from the Henry MacCraken Fellowship. G.M.R. acknowledges support from the James S. McDonnell Foundation. J.B. acknowledges support from the Alfred P. Sloan Foundation, NSF RI-1816753, NSF CAREER CIF 1845360, and the Institute for Advanced Study. E. V.-E. acknowledges support from the National Science Foundation (NSF) Materials Research Science and Engineering Center Program Grant No. DMR-1420073, and from NSF Grant No. DMS-1522767.

References

Appendix

We will use $\nabla\varphi(\bm{\theta},\bm{x})$ and $\nabla\nabla\varphi(\bm{\theta},\bm{x})$ to denote $\nabla_{\bm{\theta}}\varphi(\bm{\theta},\bm{x})$ and $\nabla_{\bm{\theta}}\nabla_{\bm{\theta}}\varphi(\bm{\theta},\bm{x})$, respectively. We will use $\nabla K(\bm{\theta},\bm{\theta}^{\prime})$ to denote $\nabla_{\bm{\theta}}K(\bm{\theta},\bm{\theta}^{\prime})$, $\nabla\nabla K(\bm{\theta},\bm{\theta}^{\prime})$ to denote $\nabla_{\bm{\theta}}\nabla_{\bm{\theta}}K(\bm{\theta},\bm{\theta}^{\prime})$, $\nabla^{\prime}\nabla K(\bm{\theta},\bm{\theta}^{\prime})$ to denote $\nabla_{\bm{\theta}^{\prime}}\nabla_{\bm{\theta}}K(\bm{\theta},\bm{\theta}^{\prime})$, and $\nabla^{\prime}\nabla^{\prime}K(\bm{\theta},\bm{\theta}^{\prime})$ to denote $\nabla_{\bm{\theta}^{\prime}}\nabla_{\bm{\theta}^{\prime}}K(\bm{\theta},\bm{\theta}^{\prime})$. We will write $V_{t}(\cdot)$ for $V(\cdot,\mu_{t})$ and $V_{\infty}(\cdot)$ for $V(\cdot,\mu_{\infty})$.

Let $D^{\prime}=\cup_{t>0}\operatorname{supp}\mu_{t}$. Under Assumption 2.5 and Proposition 2.7, $D^{\prime}$ is bounded, and we denote its diameter by $|D^{\prime}|$. We will use $C_{\varphi}$, $C_{\nabla\varphi}$ and $C_{\nabla\nabla\varphi}$ to denote the supremum of $|\varphi(\bm{\theta},\bm{x})|$, $|\nabla\varphi(\bm{\theta},\bm{x})|$ and $|\nabla\nabla\varphi(\bm{\theta},\bm{x})|$ over $\bm{\theta}\in D^{\prime}$ and $\bm{x}\in\operatorname{supp}\hat{\nu}$, which are all finite under Assumptions 2.2 and the boundedness of $D^{\prime}$. We will use $L_{\nabla\nabla\varphi}$ to denote the (uniform-in-$\bm{x}$) Lipschitz constant of $\nabla\nabla\varphi(\bm{\theta},\bm{x})$ in $\bm{\theta}$, which is also finite under Assumption 2.2.

The following notations will be used in Appendix D.2: Assuming that $D$ is Euclidean (under Assumption 2.2), let $\mathcal{V}(D)$ denote the space of random vector fields on $D$. It becomes a Hilbert space once equipped with the inner product

where $\bm{\xi}_{1}$, $\bm{\xi}_{2}$ denotes two random vector fields in $\mathcal{V}(D)$. This inner product gives rise to the norm

For each $t$, we define $\bm{b}_{t}\in\mathcal{V}(D)$ as

which depends on the random measure $\omega_{0}$. We define two linear operators, $\mathcal{A}_{t}^{(K)}$ and $\mathcal{A}_{t}^{(V)}$ on $\mathcal{V}(D)$, as

for $\bm{\xi}\in\mathcal{V}(D)$. Under Assumption 2.5, we also define $\bm{b}_{\infty}$, $\mathcal{A}_{\infty}^{(K)}$, and $\mathcal{A}_{\infty}^{(V)}$ similarly by replacing $\bm{\Theta}_{t}(\cdot)$ with $\bm{\Theta}_{\infty}(\cdot)$.

Let $\mathcal{W}_{n}(\Omega)$ denote the space of random functions on $\Omega$. It becomes a Hilbert space once equipped with the inner product

which maps a vector field $\bm{\xi}\in\mathcal{V}(D)$ back into $\mathcal{W}_{n}(\Omega)$.

Appendix B Long-Time Properties of the Mean-Field Gradient Flow

Proof of Proposition 2.7: The compactness of $\cup_{t\geq 0}\operatorname{supp}{\mu_{t}}$ follows from (20) and the compactness of $\operatorname{supp}{\mu_{0}}$ assumed in Assumption 2.3. $\mu_{t}\rightharpoonup\mu_{\infty}$ follows from (13) and (20).

Under Assumption 2.5, $\bm{\Theta}_{\infty}$ is a local minimizer of the energy $\mathcal{E}$ defined in (18). Consider a local perturbation $\epsilon\bm{\Theta}_{\Delta}$ to $\bm{\Theta}$. The energy value after the perturbation is

Under Assumptions 2.2, using Taylor expansion, we have

Since $\bm{\Theta}_{\Delta}$ is arbitrary can $\epsilon$ can be taken arbitrarily small, we see that for $\bm{\Theta}_{\infty}$ to be a local minimizer, the first-order condition is, $\forall\bm{\theta}\in\operatorname{supp}\mu_{0}$,

and the second-order condition is, $\forall\bm{\Theta}_{\Delta}$,

Therefore, for $J$ large enough, we will have

which contradicts (76). Hence, we can conclude that $\mu_{0}$-almost surely, $\nabla\nabla V(\bm{\Theta}_{\infty}(\bm{\theta}),\mu_{\infty})$ is positive semidefinite.

Appendix C Derivations of the Dynamical Central Limit Theorem

The following derivation is an adaptation of the approach in for Vlasov interacting particle systems to our scenario. To start, $\bm{\Theta}_{t}$ and $\bm{\Theta}_{t}^{(m)}$ are governed by the following equations, respectively:

Taking the difference between the two equations in (81) and using the mean value theorem, we get

This ends the proof of Proposition 3.1. $\square$

C.2 Proof of Proposition 3.3 (Dynamical CLT - II)

Since $\bm{T}_{0}(\bm{\theta})=0$, we can use Duhamel’s principle to deduce that

where the tensor $J_{t,s}(\bm{\theta})$ is the Jacobian defined in Proposition 3.3. As a result

with $\bar{g}_{t}(\bm{x})$ and $\Gamma_{t,s}(\bm{x},\bm{x}^{\prime})$ defined in (37) and (39), respectively. This is (38). $\square$

Appendix D Long-Time Behavior of the Fluctuations

With the argument outlined in Section 3.2, what remains to be shown is that $\Gamma_{t-s}^{\infty}$ is positive-semidefinite as a Volterra kernel, according to the definition in . We will utilize the following known result:

because by assumption, $\forall\bm{\theta}\in D$, $\nabla\nabla V_{\infty}(\bm{\Theta}_{\infty}(\bm{\theta}))$ is positive semidefinite, and hence $e^{-t\nabla\nabla V_{\infty}(\bm{\Theta}_{\infty}(\bm{\theta}))}$ is a positive semidefinite operator;

(2) nonincreasing: Taking derivative with respect to time,

because again, $\nabla\nabla V_{\infty}(\bm{\Theta}_{\infty}(\bm{\theta}))$ is positive semidefinite;

(3) convex: Taking one more derivative with respect to time,

Therefore, we can apply Proposition D.1 to conclude that $\Gamma_{t-s}^{\infty}$ is PSD as a Volterra kernel, and so $\int_{t_{0}}^{T}\int_{t_{0}}^{t}\langle g_{t},\Gamma^{\infty}_{t-s}g_{s}\rangle dsdt\geq 0$.

D.2 Proof of Theorem 3.5 (Unregularized case)

where here and below we denote $\fint_{0}^{t}[\cdot]\ dt=\frac{1}{t}\int_{0}^{t}[\cdot]\ dt$. As a result, to prove (46) or (49) in Theorem 3.5, it suffices to establish that

To this end, we examine (33) as an infinite-dimensional ODE. With the Hilbert space $\mathcal{V}(D)$ defined in Appendix A and $\bm{b}_{t}$, $\mathcal{A}_{t}^{(K)}$ and $\mathcal{A}_{t}^{(V)}$ defined by (61), (63) and (64), respectively, we can rewrite (33) as the following ODE on $\mathcal{V}(D)$:

Note that for all $t$, $\mathcal{A}_{t}^{(K)}$ is a positive semidefinite (PSD) operator on $\mathcal{V}(D)$, as $\forall\bm{\xi}\in\mathcal{V}(D)$,

This implies that $\fint_{0}^{T}\langle\bm{T}_{t},\mathcal{A}_{t}^{(K)}\bm{T}_{t}\rangle_{0}dt\geq 0$. Hence, to establish (99), it is sufficient to show that

To this end, we need two lemmas that are proved below in Appendices D.2.1 and D.2.2, respectively:

Assuming (47) and (48) together with Assumptions 2.2, 2.3 and 2.5, we have

Assuming (47) and (48) together with Assumptions 2.2, 2.3 and 2.5, we have

and therefore (109) is satisfied. This finishes the proof of (46) under (47) and (48) together with Assumptions 2.2, 2.3 and 2.5.

Next, we show (49) under the additional condition of Assumption 2.1. Thanks to (107) and (109), it is sufficient to establish that

Heuristically, if $\bm{T}_{\infty}:=\lim_{t\to\infty}\bm{T}_{t}$ exists, then from (102), it has to satisfy

as $\mathcal{A}_{\infty}^{(V)}=0$ (because $\nabla\nabla V(\bm{\theta},\mu_{\infty})=\int_{\Omega}\varphi(\bm{\theta},\bm{x})(f_{\infty}(\bm{x})-f_{*}(\bm{x}))\hat{\nu}(d\bm{x})=0$ under the assumption of (47)). This equation implies that

where $\left(\bm{T}_{\infty}\right)^{||}$ denotes the component of $\bm{T}_{\infty}$ in the range of $\mathcal{A}_{\infty}^{(K)}$, and $\left(\mathcal{A}_{\infty}^{(K)}\right)^{\dagger}$ denotes the Moore-Penrose pseudoinverse of $\mathcal{A}_{\infty}^{(K)}$. As a result,

Rigorously, without assuming the existence of $\bm{T}_{\infty}$, we can establish that

Assuming (47) and (48) together with Assumptions 2.2, 2.3 and 2.5, we have

This lemma is proved in D.2.3. It implies that we only need to show that

This requires us to further exploit the relationship among $\mathcal{A}_{\infty}^{(K)}$, $\bm{b}_{\infty}$ and $\bar{g}_{\infty}$. With the Hilbert space $\mathcal{W}_{L}(\Omega)$ defined in Appendix A and $\mathcal{B}_{t}$ defined by (67), we can rewrite (63) as

Similar formulas hold when we replace $t$ by $\infty$. With these relations, we see that

Lemma D.5 is proven in Appendix D.2.4 and it concludes the proof of (49) in Theorem 3.5.

To show that $\|g_{t}\|_{\hat{\nu}}$ decreases monotonically when $\mu_{0}=\mu_{\infty}$, note that in this case $\mu_{t}=\mu_{\infty}$, $\forall t\geq 0$, and so $\mathcal{A}_{t}^{(V)}=\mathcal{A}_{\infty}^{(V)}=0$, $\mathcal{A}_{t}^{(K)}=\mathcal{A}_{\infty}^{(K)}$ and $\bm{b}_{t}=\bm{b}_{\infty}$, $\forall t\geq 0$. Thus, (102) becomes

As will be shown in Lemma D.6, $\bm{b}_{\infty}$ is in the range of $\mathcal{A}_{\infty}^{(K)}$. Therefore, defining

whose solution can be written analytically as

Therefore, as $\bm{b}_{\infty}=\mathcal{B}_{\infty}\bar{g}_{\infty}$, there is

In the ERM setting, $\mathcal{A}_{\infty}^{(K)}$ is PSD with a finite number of nonzero eigenspaces. Consider a set of its orthonormal eigenfunctions that span those nonzero eigenspaces, $v_{1},...,v_{k}$, corresponding to eigenvalues $\lambda_{1},...,\lambda_{k}>0$, respectively. As $\bm{b}_{\infty}$ is in the range of $\mathcal{A}_{\infty}^{(K)}$ by Lemma D.6, we can decompose it as

for some real numbers $c_{i}$’s. Thus, we can write

which is decreasing in time. This completes the proof of Theorem 3.5. $\square$

Proof of (110): $\int_{0}^{\infty}\|\mathcal{A}_{t}^{(V)}\|_{0}dt<\infty$. By the definition of the operator norm induced by $\|\cdot\|_{0}$ on $\mathcal{V}(D)$, $\|\mathcal{A}_{t}^{(V)}\|_{0}$ is the smallest number $C_{t}$ such that $\forall\bm{\xi}$, there is

In the unregularized case, a straightforward bound of $\left|\langle\bm{\xi},\mathcal{A}_{t}^{(V)}\bm{\xi}\rangle_{0}\right|$ is

which gives us the desired bound. $\square$ Proof of (111): $\int_{0}^{\infty}\|\mathcal{A}_{\infty}^{(K)}-\mathcal{A}_{t}^{(K)}\|_{0}dt<\infty$. We have

Hence, the absolute value of the expression above is upper-bounded by

$\square$ Proof of (112): $\int_{0}^{\infty}\|\bm{b}_{t}-\bm{b}_{\infty}\|_{0}dt<\infty$. There is

By the property of $\omega_{0}$, there is

Hence, with the assumption of (48), we can conclude that

Our goal is to show that $\|\bm{T}_{t}\|_{0}$ remains bounded for all time. First note that, for all $t$, $\mathcal{A}_{t}^{(K)}$ is a positive semidefinite (PSD) operator on $\mathcal{V}(D)$ since

Second, by Assumption 2.5, for $\mu_{0}$-almost-every $\bm{\theta}\in D$, $\bm{\Theta}_{\infty}(\bm{\theta})=\lim_{t\to\infty}\bm{\Theta}_{t}(\bm{\theta})$ exists, which allows us to define $\bm{b}_{\infty}$, $\mathcal{A}_{\infty}^{(K)}$, and $\mathcal{A}_{\infty}^{(V)}$ similarly to (61), (63) and (64) by replacing $\bm{\Theta}_{t}(\cdot)$ with $\bm{\Theta}_{\infty}(\cdot)$. Since we assume that

This implies that $\mathcal{A}_{\infty}^{(V)}$ is the zero operator on $\mathcal{V}(D)$.

Third, we have the following observation:

where $\lambda_{\min}$ is the least nonzero eigenvalue of the matrix $\mathcal{M}_{\infty}$ (and hence $\lambda_{\min}^{-1}$ is the largest eigenvalue of $\mathcal{M}_{\infty}^{\dagger}$). Since

(End of the proof of Lemma D.6) $\square$

Coming back to the prof of Lemma D.3, we have shown that, as $t\to\infty$, (102) approaches the asymptotic dynamics

with $\mathcal{A}_{\infty}^{(K)}$ positive semidefinite and $\bm{b}_{\infty}$ in the range of $\mathcal{A}_{\infty}^{(K)}$. This is a stable system. Hence, the rest of the task is to examine what happens at finite time. To do so, we perform a change-of-variable with

as is defined in the proof of Lemma D.6. The dynamics of $\bm{z}_{t}$ is governed by

where $\Pi(t,s)$ is the fundamental solution (a.k.a. Green’s function) associated with the time-variant homogeneous system

Since $\mathcal{A}_{t}^{(K)}$ is positive semidefinite for all $t$, there is $\|\Pi(t,s)\|_{0}\leq 1$ for $t>s$, where with a slight abuse of notation we also use $\|\cdot\|_{0}$ for the operator norm. Hence,

Therefore, $\|\bm{z}_{t}\|_{0}$ remains bounded for all time if we can show that

we see that (172) is guaranteed by Lemmas D.2 and D.6.

This completes the proof of Lemma D.3. $\square$

By (110) in Lemma D.2 as well as Lemma D.3, we know that

Next, by (111) and (112) in Lemma D.2 as well as Lemma D.3, we know that

Let $\bm{\xi}^{||}$ denote the component of a vector field $\bm{\xi}\in\mathcal{V}(D)$ that is in the range of $\mathcal{A}_{\infty}^{(K)}$. In the ERM setting, $\mathcal{A}_{\infty}^{(K)}$ has a least nonzero eigenvalue that is positive, and hence the above implies that

Similar to (111), it can be shown that $\int_{0}^{\infty}\|\mathcal{B}_{t}-\mathcal{B}_{\infty}\|_{0}dt<\infty$. Therefore, we have

we know that when viewed as a $L$-dimensional random vector, $\bar{g}_{\infty}$ has the distribution

by the covariance of $\omega_{0}$, (32). Thus, we decompose $\bar{C}_{\infty}$ as $\bar{C}_{\infty}=\bar{C}_{\infty}^{(1)}-\bar{C}_{\infty}^{(2)}$, with

Since $\bar{C}_{\infty}$ is PSD, its square root $\left(\bar{C}_{\infty}\right)^{1/2}$ is well-defined. By the property of multivariate Gaussian, we can write

D.3 Proof of Theorem 3.6 (Under assumptions on the curvature in the long-time)

When the limiting measure $\mu_{\infty}$ does not necessarily interpolate the training data, such as in the regularized case, we have the following condition on $\bm{T}_{t}$ which guarantees that (46) holds:

(including when this limit is $+\infty$) then (46) holds.

Proof of Lemma D.7: With $\mathfrak{D}_{t}$ defined in (101), for (46) to hold, it is sufficient to show that

Since $\bm{T}_{0}=0$ and $\mathcal{A}_{t}^{(K)}$ is PSD, we see that the assumption (200) is sufficient. $\square$

Note that condition (200) is natural since we know from Proposition 2.7 that $\lim_{t\to\infty}\nabla\nabla V(\bm{\Theta}_{t}(\bm{\theta}),\mu_{t})=\nabla\nabla V(\bm{\Theta}_{\infty}(\bm{\theta}),\mu_{\infty})$ exists and is positive semidefinite $\mu_{0}$-almost surely. This lemma then allows us to prove Theorem 3.6: Proof of Theorem 3.6: Our goal is to verify (200) in order to apply Lemma D.7. We first see that

where we define, for $\bm{\xi}\in\mathcal{V}(D)$,

Hence, if we assume that $\left|\int_{D}\min\left\{\lambda_{\min}(\nabla\nabla V(\bm{\theta},\mu_{t})),0\right\}\mu_{0}(d\bm{\theta})\right|$ is small asymptotically, then what remains is to upper-bound $\|\bm{T}_{t}\|_{\sup}$. Recall from (102) that the dynamics of $\bm{T}_{t}$ is governed by

Thus, in the $\|\cdot\|_{\sup}$ norm defined above, we have

We then want to bound the growth of $\|\bm{T}_{t}\|_{\sup}$ by upper-bounding the RHS. Note that for $\bm{\xi}\in\mathcal{V}(D)$,

To bound $\|\bm{b}_{t}\|_{\sup}$, we recall that

This implies that $\forall\bm{\theta}\in\operatorname{supp}\mu_{0}$,

On the other hand, similar to (162), we have

Thus, there is $\forall\bm{\theta}\in\operatorname{supp}\mu_{0}$,

Now, using (203), we see that in order for (200) to hold, it is sufficient to have

To intuitively understand (50), note that we know from (23) in Proposition 2.7 that $\Lambda_{t}(\bm{\theta})\to 0$ $\mu_{0}$-almost surely as $t\to\infty$. Condition (50) can therefore be satisfied by having $\Lambda_{t}(\bm{\theta})$ converge to zero sufficiently fast in the regions of $D$ where it is negative, or having the measure of these regions with respect to $\mu_{0}$ converge to zero sufficiently fast, or both.

D.4 Proof of Theorem 3.7 (Regularized case)

Recall from Proposition 3.3 that the dynamics of $g_{t}$ is governed by

with $J_{t,s}$ being the Jacobian of the flow $\bm{\Theta}_{t}$.

In the ERM setting, $\operatorname{supp}\hat{\nu}$ is singular, thus we have $\hat{\nu}(d\bm{x})=n^{-1}\sum_{l=1}^{n}\delta_{\bm{x}_{l}}(d\bm{x})$, where $n$ is the total number of training data points. We define $\mathcal{W}_{n}(\Omega)$ together with the inner product $\langle\cdot,\cdot\rangle_{\hat{\nu},0}$ and the norm $\|\cdot\|_{\hat{\nu},0}$ as in Appendix A. We will also continue to consider $g_{t}$ and $\bar{g}_{t}$ equivalently as $n$-dimensional vectors,

respectively. Thus, $\Gamma_{t,s}$ can also be represented by the $n\times n$ matrix

Under such an abuse of notations, we can simplify (221) into

where for simplicity, we write $V_{t}(\cdot)$ for $V(\cdot,\mu_{t})$ and $V_{\infty}(\cdot)$ for $V(\cdot,\mu_{\infty})$. Then the heuristic argument outlined in Section 3.2 before Theorem 3.4 amounts to rewriting (225) as

and then arguing that 1) $\Gamma^{\infty}$ is a nonnegative convolution-type Volterra kernel, and 2) the second term on the RHS is small. Rigorously, we need to introduce an extra level of complication: for every $t_{0}>0$, we can rewrite (225) into

Then, for any $T>t_{0}$, by multiplying $g_{t}$ and integrating from $t_{0}$ to $T$, we get

Then firstly, the second term on the LHS is nonnegative because of the nonnegativity of $\Gamma_{t}^{\infty}$ as a convolution-type Volterra kernel, as proven in Appendix D.1.

Therefore, putting everything together, we have

and hence, using $\fint_{a}^{b}\cdot\ dt$ to denote the averaged integral $\frac{1}{b-a}\int_{a}^{b}\cdot\ dt$,

Under all assumptions in Theorem 3.7 except for (52) being replaced by a weaker condition,

We will prove in Appendix D.4.2 that (52) indeed implies (238).

The lemma will be proved in Appendix D.4.1, and let us first proceed with the proof of the theorem assuming this lemma. Suppose for contradiction that (226) does not hold, meaning that

for some $\epsilon>0$. We will select a pair of $t_{0}$ and $T$ for which the inequality (237) cannot be satisfied. Firstly, by the convergence of $\bar{g}_{t}$ to $\bar{g}_{\infty}$, $\exists t_{a}>0$ such that $\forall t_{1},t_{2}>t_{a}$,

Secondly, by our assumption (241) and the first part of Lemma D.8, $\exists t_{0}>t_{a}$ such that both

are satisfied. In particular, (243) implies

By the second part of Lemma D.8, $\exists t_{b}>t_{0}$ such that $\forall T>t_{b}$,

By our assumption (241), we can choose a $T>t_{b}$ such that

we can assume without loss of generality that $\frac{T}{t_{0}}$ is large enough so that

Thus, back to the inequality (237), the LHS is strictly lower-bounded by

whereas the RHS is strictly upper-bounded by

This gives contradiction and we are done with the proof of Theorem 3.7. $\square$

It remains to prove Lemma D.8. To do so we will need an auxiliary result, that we state and prove first:

Let $\Delta\Gamma_{t,s}:=\Gamma_{t,s}-\Gamma_{t-s}^{\infty}$. If $\nabla\nabla V$ is uniformly positive definite with eigenvalues lower-bounded by $\lambda$, then there exists constants $C$ and $C^{\prime}$ whose values depend on $|D^{\prime}|$, $C_{\varphi}$, $C_{\nabla\varphi}$, $C_{\nabla\nabla\varphi}$, and $L_{\nabla\nabla\varphi}$ such that

where $\Delta\bm{\Theta}_{t}(\bm{\theta})=\bm{\Theta}_{t}(\bm{\theta})-\bm{\Theta}_{\infty}(\bm{\theta})$.

We will first try to bound $\bm{\xi}_{t}(\bm{\theta})-\bm{\xi}^{\prime}_{t}(\bm{\theta})$ as a function of $\eta$. Define $\Delta\bm{\xi}_{t}(\bm{\theta})=\bm{\xi}_{t}(\bm{\theta})-\bm{\xi}^{\prime}_{t}(\bm{\theta})$. Then

Since $\nabla\nabla V_{\infty}(\bm{\Theta}_{\infty}(\bm{\theta}))-\lambda I_{d}$ is positive semidefinite, we first have

To prepare for an application of Gronwall’s inequality, we introduce a change-of-variable by defining, for $r\in[s,t]$,

Then we can rewrite the equation above as

or, back in the original variable that we are interested in,

and, since $\nabla\nabla V_{r}(\bm{\theta})=\int_{\Omega}\nabla\nabla\varphi(\bm{\theta},\bm{x})(f_{r}(\bm{x})-f_{*}(\bm{x}))\hat{\nu}(d\bm{x})$ and $\nabla\nabla V_{\infty}(\bm{\theta})=\int_{\Omega}\nabla\nabla\varphi(\bm{\theta},\bm{x})(f_{\infty}(\bm{x})-f_{*}(\bm{x}))\hat{\nu}(d\bm{x})$,

where we use $\|f\|_{\hat{\nu},\infty}$ to denote $\sup_{\bm{x}\in\operatorname{supp}\hat{\nu}}|f(\bm{x})|$ and we defined $\Delta f_{t}=f_{t}-f_{\infty}$.

Therefore, using $C_{0}$, $C_{1}$, etc. to represent constants that depend on $C_{\varphi}$, $C_{\nabla\varphi}$, $C_{\nabla\nabla\varphi}$, $C_{\nabla\nabla\varphi}$ and $L_{\nabla\nabla\varphi}$, we have

Note that $\|\Delta f_{r}\|_{\hat{\nu},\infty}$ can be further upper-bounded by $C_{\varphi}\int_{D}|\Delta\bm{\Theta}_{r}(\bm{\theta})|\mu_{0}(d\bm{\theta})$. Furthermore, defining

(End of the proof of Lemma D.9.) $\square$

Proof of Lemma D.8: Lemma D.9 entails that, $\exists C,C^{\prime}>0$ such that

where for the last inequality, we assume that $|D^{\prime}|\geq 1$ (or, to accommodate the more general case, just replace $|D^{\prime}|$ by $\max\{|D^{\prime}|,1\}$).

To prove Lemma D.8, the first goal is to show

By our assumption, the RHS is finite for $t_{0}>0$. Hence, by taking $t_{0}$ large enough, the value of $\int_{t_{0}}^{\infty}\int_{t_{0}}^{t}\|\Delta\Gamma_{t,s}\|_{\hat{\nu}}^{2}dsdt$ can be made arbitrarily close to zero.

The second goal is to show that $\forall t_{0}>0$,

because $\forall\eta\in\mathcal{W}_{L}(\Omega)$, there is

and $\mathcal{M}_{\infty}$ is defined as $\mathcal{M}_{\infty}:=\mathcal{B}_{\infty}^{\intercal}\mathcal{B}_{\infty}$, or concretely, for $\eta\in\mathcal{W}_{L}(\omega)$,

In the ERM setting, $\mathcal{M}_{\infty}$ is effectively an $L\times L$ matrix. Thus,

This concludes the proof of Lemma D.8. $\square$

Below, we will illustrate the assumption (238)

in Theorem 3.7 by giving examples that satisfy this condition.

First, consider an example where $\exists\kappa>0,\alpha>1$ such that $\forall\bm{\theta}\in\operatorname{supp}\mu_{0}$ and $\forall~{}t>0$,

that is, all characteristic flows share a uniform asymptotic convergence rate on the order of $t^{-\alpha}$. Then $\forall\bm{\theta}\in\operatorname{supp}\mu_{0}$,

which is finite as long as $\alpha>\frac{3}{2}$. Thus,

If $\exists\kappa>0,\alpha>\frac{3}{2}$ such that $\forall\bm{\theta}\in\operatorname{supp}\mu_{0}$ and $\forall t\geq 0$,

such that $\forall\bm{\theta}\in\operatorname{supp}\mu_{0}$,

Appendix E Properties of the Minimizers of the Regularized Loss

First, under Assumption 2.1, i.e. in the shallow neural networks setting, define

We prove Proposition 3.9, which we extend into:

Under Assumptions 2.1, 2.2, and 3.8, the minimizers of the loss $\mathcal{L}(\mu)$ defined in (4) are all in the form

where $c_{\lambda}\geq 0$ and $\hat{\mu}_{\pm}\in\mathcal{P}(\hat{D})$ satisfy

In addition, the constant $c_{\lambda}$ is unique and positive if $\hat{F}(\bm{z})$ is not identically zero on $\hat{D}$, the closure of the supports of $\hat{\mu}_{\pm}$ are disjoint (i.e. $\overline{\operatorname{supp}\hat{\mu}_{+}}\cap\overline{\operatorname{supp}\hat{\mu}_{-}}=\emptyset$), and the function

is the same for all minimizers and satisfies

where $\hat{K}_{M}=\max_{\bm{z}\in\hat{D}}\left\|\hat{\varphi}(\bm{z},\cdot)\right\|_{\hat{\nu}}^{2}=\max_{\bm{z}\in\hat{D}}\hat{K}(\bm{z},\bm{z})$.

Note that the proposition automatically implies that $\gamma_{1}(f_{\lambda})\leq\gamma_{1}(f_{*})<\infty$. It also implies that

with equality on the support of $\mu$ and where $\bar{V}$ is the expectation of the left hand side with respect to $\mu(dc,d\bm{z})$. Minimizing the left hand side of (309) over $c$ at fixed $\bm{z}$, we deduce that

If we insert this equality back in $c(\bm{z})\hat{V}(\bm{z})+\frac{1}{2}\lambda|c(\bm{z})|^{2}=\bar{V}$, we deduce that $|c(\bm{z})|=c_{\lambda}$, with the constant $c_{\lambda}$ related to $\bar{V}$ as

and furthermore, $\forall\bm{z}\in\operatorname{supp}{\hat{\mu}}$,

These considerations imply that the minimizer must be of the form (304), and if we combine (310) and (312) and evaluate the minimum on $c$ explicitly we deduce that $\hat{\mu}_{\pm}$ and $c_{\lambda}$ must satisfy the equations in (305). It is also clear from (305) that we must have $\overline{\operatorname{supp}\hat{\mu}_{+}}\cap\overline{\operatorname{supp}\hat{\mu}_{-}}=\emptyset$: indeed if there was a point $\bm{z}\in\overline{\operatorname{supp}\hat{\mu}_{+}}\cap\overline{\operatorname{supp}\hat{\mu}_{-}}$, then at that point $\hat{V}(\bm{z})$ would be discontinuous, which is not possible since this function is continuously differentiable for any $\mu$ by our assumptions on $\hat{\varphi}$. Finally, to show that we must have that $c_{\lambda}>0$ if $F(\bm{z})$ is not identically zero on $\hat{D}$, note that if $c_{\lambda}=0$, (309) reduces to

which can only be satisfied if $\hat{F}(\bm{z})=0$.

To show that $c_{\lambda}$ and the function in (306) are unique, let $\mu_{\lambda}$ and $\mu^{\prime}_{\lambda}$ be two different minimizers and consider

Let us evaluate the loss on $a\mu_{\lambda}+(1-a)\mu^{\prime}_{\lambda}\in\mathcal{P}(D)$ with $a\in$. By convexity of $\mathcal{E}_{\lambda}$ we have

Since $a\mu_{\lambda}+(1-a)\mu^{\prime}_{\lambda}$ cannot have a lower loss than this minimum, we must have equality in (316), which reduces to

where $c_{\lambda}$ and $c_{\lambda}^{\prime}$ are associated with $\mu_{\lambda}$ and $\mu_{\lambda}^{\prime}$, respectively. Clearly these equations can only be fulfilled for all $a\in$ if $c_{\lambda}=c_{\lambda}^{\prime}$ and $f_{\lambda}=f_{\lambda}^{\prime}$ ${\hat{\nu}}$-a.e. on $\Omega$.

To establish (307), notice that if $\mu_{\lambda}$ is a minimizer and $f_{\lambda}$ is given by (306), then we can derive from (313) that