Mean-field Analysis of Piecewise Linear Solutions for Wide ReLU Networks

Introduction

Neural networks are the key ingredient behind many recent advances in machine learning. They achieve state-of-the-art performance on various practical tasks, such as image classification [HZRS16] and synthesis [BDS19], natural language processing [VSP+17] and reinforcement learning [SHM+16]. However, these results would not be possible without computational advances which enabled the training of highly overparameterized models with billions of weights. Such complex networks are capable of extracting more sophisticated patterns from the data than their less parameter-heavy counterparts. Nonetheless, in the view of classical learning theory, models with a large number of parameters are prone to over-fitting [VLS11]. Contrary to the conventional statistical wisdom, overparameterization turns out to be a rather desirable property for neural networks. This was even observed in a classical paper by [Bar98], which demonstrated that in the overparameterized setting, the size of the network is less important than the magnitude of the weights. More recently, phenomena such as double descent [BHMM19, SGd+19, NKB+20] and benign overfitting [BLLT20, LZG21, BMR21] suggest that understanding the generalization properties of overparameterized models lies beyond the scope of the usual control of capacity via the size of the parameter set [NTS15].

One way to explain the generalization capability of large neural networks lies in characterizing the properties of solutions found by stochastic gradient descent (SGD). In other words, the question is whether the optimization procedure is implicitly selective, i.e., it finds the functionally simple solutions that exhibit superior generalization ability in comparison to other candidates with roughly the same value of the empirical risk. For instance, [CB20] consider shallow networks minimizing the logistic loss, and show that SGD converges to a max-margin classifier on a certain functional space endowed with the variation norm. In the machine learning literature, it has been suggested that large margin classifiers inherently exhibit better performance on unseen data [BMR21, CV95].

In this work, we develop a novel approach towards understanding the implicit bias of gradient descent methods applied to overparameterized neural networks. In particular, we focus on the following key questions:

Once stochastic gradient descent has converged, how does the distribution of the weights of the neural network look like? What functional properties of the resulting solution are induced by this stationary distribution? Can we quantitatively characterize the trade-off between the complexity of the solution and the size of the training data in the overparameterized regime?

To answer these questions, we consider training a wide two-layer ReLU (rectified linear unit) network for univariate regression, and we focus on the mean-field regime [MMN18, RVE18, CB18, SS20]. In this regime, the idea is that, as the number of neurons of the network grows, the weights obtained via SGD are close to i.i.d. samples coming from the solution of a certain Wasserstein gradient flow. As a consequence, the output of the neural network approaches the following quantity:

Here, $x$ is the input, $\sigma^{*}$ denotes the activation function, and $\rho$ is the solution of the Wasserstein gradient flow minimizing the free energy

\beta=+\infty, while in (c) we have a sufficiently large regularization coefficient, which does not allow the network to fit the training data perfectly. Note that the piecewise linear solution exhibits tangent changes also at points different from the training data. Furthermore, the number of “knot” points may differ from the minimum required to fit the data: for instance, in (a) the minimum amount of tangent changes is $1$, but the solution has two of them. A number of works has exploited this connection to provide a rigorous justification to various phenomena attributed to neural networks. [MMN18, MMM19] give global convergence guarantees for two-layer networks by studying the energy dissipation along the trajectory of the flow. The paper by [CB18] takes a different route and exploits a lifting property enabled by a certain type of initialization and regularization, and [JMM20] put forward an argument based on displacement convexity. [NP20] and [AOY19] tackle the multi-layer case, and, in particular, [NP20] establish convergence guarantees for a three-layer network. [FLYZ21] introduce a mean-field dynamics capturing the evolution of the features (instead of the network parameters) and show global convergence of ResNet type of architectures. [SM20] prove two properties commonly observed in practice (see e.g. [GIP+19, DVSH18, KWL+19]), namely dropout stability and mode connectivity, for multi-layer networks trained under the mean-field regime. [DBDFS20] consider different scalings of the step size of SGD, and identify two regimes under which different mean-field limits are obtained. [WTS+19] show that the gradient flow for unregularized objectives forces the neurons of a two-layer ReLU network to concentrate around a subset of the training data points.

In this paper, we take a mean-field view to show that SGD is biased towards functionally simple solutions, namely, piecewise linear functions. Our idea is to analyze the stationary distribution $\rho^{*}$ minimizing the free energy (1.1). We show that, in the low temperature regime ($\beta\rightarrow\infty$), the estimator’s curvature vanishes everywhere except for a certain cluster set. More precisely, for each interval between two consecutive training inputs, aside for a set of small measure, the second derivative vanishes, i.e.,

Furthermore, we provide a characterization of the cluster set and show that its measure vanishes while it concentrates around at most 3 points per interval. Ultimately, this analysis guarantees that, in the regime of decreasing temperature (corresponding to a small noise injected in the gradient updates), the solution found by SGD is piecewise linear. Our main contribution can be summarized in the following informal statement:

Under the low temperature regime, i.e., $\beta\rightarrow\infty$, the estimator obtained by training a two-layer ReLU network via noisy-SGD converges to a piecewise linear solution. Furthermore, the number of “knot” points – i.e., points at which distinct linear pieces connect – between two consecutive training inputs is at most 3.

Let us remark on a few important points. In the overparameterized regime, the number of neurons $N$ is significantly larger than the number of training samples $M$, i.e., $N\gg M$. The output of the two-layer ReLU network is a linear combination of $N$ ReLU units, hence the function implemented by the network is clearly piecewise linear with $\mathcal{O}(N)$ knot points. Here, we show that the number of knot points is actually $\mathcal{O}(M)\ll\mathcal{O}(N)$. Our analysis applies for both constant ($\lambda\to\bar{\lambda}>0$) and vanishing ($\lambda\rightarrow 0$) regularization, and it does not require a specific form for the initialization of the parameters of the networks (as long as some mild technical conditions are satisfied).

Organization of the paper. The rest of the paper is organized as follows. In Section 2, we review the related work and a more detailed comparison is deferred to Section 8. In Section 3, we provide some preliminaries, including a background on the mean-field analysis in Section 3.1. Our main results are stated in Section 4 and proved in Section 5. In Section 6, we provide an example of a dataset for which the estimator found by SGD has a knot at a location different from the training inputs. We validate our findings with numerical simulations for different regression tasks in Section 7. We conclude and discuss some future directions in Section 9. Some of the technical lemmas and the corresponding proofs are deferred to Appendix A.

Related Work

The line of works [WTS+19, JM20] shows that, in the lazy training regime [COB19, JGH18] and for a uniform initialization, SGD converges to a cubic spline interpolating the data. Furthermore, for multivariate regression in the lazy training regime, [JM20] proved that the optimization procedure is biased towards solutions minimizing the 2-norm of the Radon transform of the fractional Laplacian. Similar results (although without the connection to the training dynamics) are obtained in [SESS19, OWSS20], which analyze the solutions with zero loss and minimum norm of the parameters. [EP21] develop a convex analytic framework to explain the bias towards simple solutions. In particular, an explicit characterization of the minimizer is provided, which implies that an optimal set of parameters yields linear spline interpolation for regression problems involving one dimensional or rank-one data. [CFW+21] show that, for overparameterized models, the lower degree spherical harmonics are easier to learn. This observation comes from the fact that, in the lazy training regime, the convergence occurs faster along the directions given by the top eigenfunctions of the neural tangent kernel. Classification with linear networks on separable data is considered in [SHN+18], where it is shown that gradient descent converges to the max-margin solution. This max-margin behavior is demonstrated in [CB20] for non-linear wide two-layer networks using a mean-field analysis. In particular, in the mean-field regime, optimizing the logistic loss is equivalent to finding the max-margin classifier in a certain functional space. The paper by [ZXLM20] focuses on the lazy training regime, and it shows that the optimization procedure finds a solution that fits the data perfectly and is closest to the starting point of the dynamics in terms of Euclidean distance in the parameter space. [WZBG21] characterize the directional bias of GD and SGD in the case of moderate (but annealing) learning rate.

The behavior of SGD with label noise near the zero-loss manifold is studied in [BGVV20]. Here, it is shown that the training algorithm implicitly optimizes an auxiliary objective, namely, the sum of squared norms of the gradients evaluated at each training sample. This allows the authors of [BGVV20] to show that SGD with label noise for a two-layer ReLU network with skip-connections is biased towards a piecewise linear solution. In particular, this piecewise linear solution has the minimum number of tangent changes required to fit the data. [WTS+19] consider the Wasserstein gradient flow on a certain space of reduced parameters (in polar coordinates), and show that the points where the solution changes tangent are concentrated around a subset of training examples. A trade-off between the scale of the initialization and the training regime is also provided in [WTS+19, SPD+20]. [MBG18] prove that the gradient flow enforces the weight vectors to concentrate at a small number of directions determined by the input data. Through the lens of spline theory, [PN20b] explain that a number of best practices used in deep learning, such as weight decay and path-norm, are connected to the ReLU activation and its smooth counterparts. [NLB+19] suggest a novel complexity measure for neural networks that provides a tighter generalization for the case of ReLU activation.

Preliminaries

where the integral is taken over the support of $\rho$. For the forthcoming analysis, a certain regularity is required for the weight distribution $\rho$. In particular, the weight distribution is restricted to a set of admissible densities

stands for the differential entropy of $\rho$, which is equal to $-\infty$ if the distribution $\rho$ is singular. In this view, the distributional dynamics (3.1) is the Wasserstein gradient flow that minimizes the free energy

over the set of admissible densities $\mathcal{K}$. Furthermore, this free energy has a unique minimizer and the solution of (3.1) converges to it as $t\rightarrow\infty$:

The unique minimizer $\rho^{*}_{\sigma^{*}}$ is absolutely continuous, and it has the Gibbs form

where $Z_{\sigma^{*}}(\beta,\lambda)$ is the normalization constant, also referred to as partition function.

2 Approximation of the ReLU Activation

Let us elaborate on the properties which $\sigma^{*}$ should satisfy so that the results of Section 3.1 hold. First, the distributional dynamic (3.1) is known to be well-defined for a smooth and bounded potential $\Psi_{\lambda}$. In particular, it suffices to choose a bounded, Lipschitz $\sigma^{*}$ with Lipschitz gradient, see assumptions A2-A3 in [MMN18]. Furthermore, the minimizer of the free energy (3.5) exists and has a Gibbs form even for non-smooth potentials and, in particular, it suffices that $\sigma^{*}$ is bounded and Lipschitz (this allows the first derivative to be discontinuous), see Lemmas 10.2-10.4 in [MMN18].

In the case of a ReLU activation, the corresponding $\sigma^{*}$ has the following form

which does not satisfy some of the aforementioned conditions. The first salient problem is the lack of continuity of the derivative at zero. This issue can be dealt with by considering a soft-plus activation with scale $\tau$:

Notice that, as $\tau$ grows large, we have that $(\cdot)_{\tau}\rightarrow(\cdot)_{+}$. Another issue is that the function $\sigma^{*}(x,{\bm{\theta}})$ is not Lipschitz in the parameters ${\bm{\theta}}$, and it is unbounded. This problem can be solved by an appropriate truncation applied to the parameter $a$ of the activation. The truncation should be Lipschitz and smooth for the dynamics to be well-defined.

We define the smooth $m$-truncation $(\cdot)_{+}^{m}$ of the ReLU activation as

Finally, we define the smooth $m$-truncation $(\cdot)_{\tau}^{m}$ of the softplus activation as

Notice that, for $\tau\geq 1$, the soft-plus activation can be sandwiched as follows:

In order to establish the continuity of a certain limit and smoothness properties, we also pick $\phi_{\tau,m}$ such that the smooth $m$-truncation of soft-plus activation satisfies a similar bound:

At this point, we remark that the activation $({\bm{\theta}},x)\mapsto a^{\tau,m}(w^{\tau,m}x+b)^{m}_{\tau}$ satisfies all the conditions necessary for the results of Section 3.1 to hold. In what follows, we will also use the activation $({\bm{\theta}},x)\mapsto a^{m}(w^{m}x+b)^{m}_{\tau}$ as an auxiliary object. This map is not smooth, but it satisfies all the assumptions required for the existence of a free energy minimizer $\rho^{*}_{\sigma^{*}}$. We also note that the truncation of the parameter $w$ might seem unnatural (we are truncating the ReLU activation anyway), but it simplifies our analysis. In particular, it allows us to establish a connection between the derivatives (w.r.t. the input $x$) of the predictor implemented by the solution of the flow (3.1) and the same quantity evaluated on the minimizer, as $t$ grows large.

We will use the following notation for the values of the risks corresponding to different activations

Here, $R^{\tau,m}_{i}$ and $R^{m}_{i}$ represent the rescaled error on the $i$-th training sample, and $R^{\tau,m}$ and $R^{m}$ are the standard expected square losses. In this way, we can write the Gibbs minimizers in a compact form, namely,

where $Z_{\tau,m}(\beta,\lambda)$ and $Z_{m}(\beta,\lambda)$ denote the partition functions.

Main Results

Before presenting the main results, let us introduce the notion of a cluster set. This set allows us to identify the locations of the knot points of an estimator function that is implemented by the neural network. In particular, we consider the second derivative of the predictor evaluated at the Gibbs distribution with activation $({\bm{\theta}},x)\mapsto a^{\tau,m}(w^{\tau,m}x+b)^{m}_{\tau}$, for large $\tau$, i.e.,

Then, the cluster set is associated to the inputs on which the quantity (4.1) might grow unbounded in absolute value, in the low temperature regime ($\beta^{-1}\rightarrow 0$). Intuitively, this indicates that on some points of the cluster set, the tangent of the predictor changes abruptly, resulting in “knots”. We denote the cluster set by $\Omega(m,\beta,\lambda)$, and we define it below.

Let $\mathcal{I}$ be the set of prediction intervals, i.e.,

where $L>\max\{|x_{1}|,\cdots,|x_{M}|\}$ is any fixed positive constant independent of $(\tau,m,\beta,\lambda)$. For each $I_{j}:=[x_{j},x_{j+1}]\in\mathcal{I}$, the intersection of the cluster set with the prediction interval $I_{j}$ is denoted by $\overline{\Omega}_{j}(m,\beta,\lambda)$, i.e.,

Thus, in order to define the cluster set $\Omega(m,\beta,\lambda)$, it suffices to give the definition of $\overline{\Omega}_{j}(m,\beta,\lambda)$. To do so, consider the second-degree polynomials $f^{j}(x)$ and $f_{j}(x)$ given by

Here, if the summation set is empty (e.g., for $A_{0}$), the corresponding coefficient is equal to zero. Then, the set $\overline{\Omega}_{j}(m,\beta,\lambda)$ is defined as the union of the non-positive sets of the second-degree polynomials $f^{j}(x)$ and $f_{j}(x)$:

We now provide an informal explanation on how the non-positive sets of the second-degree polynomials $f^{j}(x)$ and $f_{j}(x)$ come into play. A central object of interest in our analysis is the second derivative of the estimator implemented by the neural network, and our strategy is to bound its magnitude by a particular Gaussian-like integral. This integral does not diverge as long as the corresponding covariance matrix is non-degenerate, i.e., it has strictly positive eigenvalues. In this view, the non-positive sets of the polynomials $f^{j}(x)$ and $f_{j}(x)$ precisely characterize the inputs $x$ for which this covariance matrix is degenerate. Hence, for such inputs $x$, this upper bound on the second derivative of the estimator diverges, which implies that the predictor may have a “knot”.

Since $f^{j}(x)$ and $f_{j}(x)$ are second-degree polynomials, the set $\overline{\Omega}_{j}(m,\beta,\lambda)$ can be always written as the union of at most $3$ intervals. Moreover, $\overline{\Omega}_{j}(m,\beta,\lambda)$ depends only on the errors of the estimator at the training points and on the penalty parameter $\lambda$. Thus, if one has access to the value of the errors at each training point for the optimal estimator, i.e., $R^{m}_{i}(\rho^{*}_{m})$, an explicit expression for the cluster set can be readily obtained. Figure 3 shows three different configurations of the polynomials $f^{j}(x)$ and $f_{j}(x)$, together with the corresponding cluster set.

The size of the set $\overline{\Omega}_{j}(m,\beta,\lambda)$ can be controlled explicitly as a function of the parameters $(m,\beta,\lambda)$. More formally, in Lemma 5.3, we show that the Lebesgue measure of $\overline{\Omega}_{j}(m,\beta,\lambda)$ can be upper bounded as

where $C>0$ denotes a numerical constant independent of $(\tau,m,\beta,\lambda)$ and we have made the following assumption:

$\tau\geq 1$, \beta\geq\max\Big{\{}C_{1},\frac{1}{\lambda},\frac{1}{\lambda}\log\frac{1}{\lambda}\Big{\}}, $m>C_{2}$ and $\lambda<C_{3}$ for some numerical constants $C_{1},C_{2},C_{3}>0$.

In particular, (4.7) implies that the cluster set vanishes as $\beta\rightarrow\infty$ and $m=e^{\Theta(\beta)}$. Therefore, as $\overline{\Omega}_{j}(m,\beta,\lambda)$ is the union of at most $3$ intervals, the cluster set concentrates on at most 3 points per prediction interval.

We note that our use of A1 throughout the sequel is with the flexibility of $C_{1}$, $C_{2}$, and $C_{3}$ in mind; we are interested in the behavior as $m$ and $\beta$ grow large, so we permit liberty in the determination of the constants implying the formal statements we intend to show.

A key step of our analysis (cf. Theorem 1) consists in showing that, outside the cluster set, the absolute value of the second derivative vanishes. Our bound on this absolute value is connected to the speed of decay to zero of the polynomials $f^{j}(x)$ and $f_{j}(x)$, as the input $x$ approaches the cluster set. In order to establish a quantitative bound for such a decay, we introduce an auxiliary quantity, namely, a critical point, that is associated to each input point outside of the cluster set. Given the polynomial $f^{j}(\cdot)$ and the input $x\in I_{j}\setminus\Omega^{j}(m,\beta,\lambda)$, the critical point $x_{c}$ associated to $x$ is defined below.

Figure 4 provides a visualization of the critical point associated to $x$ for several configurations of $f^{j}$. For the polynomial $f_{j}(\cdot)$ and an input $x\in I_{j}\setminus\Omega_{j}(m,\beta,\lambda)$, the critical point $\bar{x}_{c}$ is defined in a similar fashion. In this view, we show in Lemma 5.5 that the following lower bounds on $f^{j},f_{j}$ hold for $x\in I_{j}\setminus\overline{\Omega}(m,\beta,\lambda)$,

The coefficients $\gamma_{1},\gamma_{2},\gamma_{3},\gamma_{4}>0$ satisfy the following condition: either $\gamma_{1}>\varepsilon$ or $\gamma_{2}>\varepsilon$, and either $\gamma_{3}>\varepsilon$ or $\gamma_{4}>\varepsilon$, where $\varepsilon>0$ is a numerical constant independent of the choice of $(m,\beta,\lambda)$.

At this point, we are ready to state our upper bound on the second derivative outside the cluster set.

Assume that condition A1 is satisfied and that $m>e^{K_{1}\beta}$ for some numerical constant $K_{1}>0$ independent of $(\tau,m,\beta,\lambda)$. Then, for each $x\in I_{j}\setminus\overline{\Omega}_{j}(m,\beta,\lambda)$, the following upper bound on the second derivative holds

where the coefficient $\bar{C}^{j}(x)$ is defined as

with $C^{j}(x)$ and $C_{j}(x)$ given by (4.8). Furthermore, the following upper-bound on the size of the cluster set holds

for some numerical constant $K_{2}>0$ independent of $(\tau,m,\beta,\lambda)$.

Some remarks are in order. First, the inequality (4.9) shows that, in the low temperature regime, the curvature vanishes outside the cluster set, and it also provides a decay rate. Second, we will upper bound the measure of the cluster set as in (4.7), thus the condition $m>e^{K_{1}\beta}$ ensures that the upper bound (4.11) holds. Finally, the presence of the coefficient $\bar{C}^{j}(x)$ is due to the fact that the second derivative can grow unbounded for points approaching the cluster set. Let us highlight that this growth is solely dictated by the distance to the cluster set, and it does not depend on $(m,\beta,\lambda)$. In fact, (4.8) holds, where one of the coefficients in $\{\gamma_{1},\gamma_{2}\}$ and in $\{\gamma_{3},\gamma_{4}\}$ is lower bounded by a strictly positive constant independent of $(m,\beta,\lambda)$.

Figure 5 provides some examples of piecewise linear solutions: (a) and (b) are admissible (in the sense of Definition 4.2), while (c) is not admissible, since it has two knots in the interior of the prediction interval and one located at the right endpoint. As mentioned before, the location of the knot points is associated with the limiting behaviour of the corresponding polynomials $f^{j}(x)$ and $f_{j}(x)$. For instance, consider the prediction interval $[x_{2},x_{3}]\in\mathcal{I}$. Then, the configuration of Figure 5a corresponds to the case described in Figure 3a. In fact, $f^{j}$ has a negative leading coefficient, and its roots are converging to the end points of the interval. Moreover, $f_{j}$ has positive curvature and the minimizer is located inside the interval. The same parallel can be drawn between Figure 5b and Figure 3c. Furthermore, one can verify that the situation described in Figure 5c cannot be achieved for any configuration of $f^{j}(x)$ and $f_{j}(x)$.

We are now ready to state our result concerning the structure of the function obtained from the Gibbs distribution $\rho^{*}_{\tau,m}$.

Assume that condition A1 is satisfied and that $m>e^{K_{1}\beta}$, where $K_{1}>0$ is a constant independent of $(\tau,m,\beta,\lambda)$. Then, given a set of prediction intervals $\mathcal{I}$, there exists a family of admissible piecewise linear solutions $\{f_{m,\beta,\lambda}\}$ as per Definition 4.2, such that, for any $I\in\mathcal{I}$ and $x\in I$, the following convergence result holds

The proof of Theorem 2 crucially relies on the fact that the second moment of $\rho^{*}_{\tau,m}$ is uniformly bounded along the sequence $\beta\lambda^{7/4}\rightarrow\infty$. In fact, the uniform bound on the second moment implies that the first derivatives of the predictors w.r.t. the input are uniformly bounded (even for points inside the cluster set), and therefore the sequence of predictors is equi-Lipschitz. This, in particular, allows us to show that the limit is well-behaved, as function changes can be controlled via Lipschitz bounds.

Let us clarify that Theorem 2 does not establish the uniqueness of the limit in $(m,\beta,\lambda)$, i.e., that the limiting piecewise linear function is the same regardless of the subsequence. Our numerical results reported in Figures 1, 6b, 7 and 8 suggest that the limit is unique. However, a typical line of argument (see e.g. [JKO98]) would require the lower-semicontinuity of the free energy (which does not hold for $m=\infty$). Furthermore, even the uniqueness of the minimizer for $\beta=\infty$ remains unclear in our setup. Nevertheless, let us point out that the sequence $\{\rho^{*}_{\tau,m}\}$ is tight, since the second moments are uniformly bounded by Lemma A.6, and Proposition 2.3 in [HRŠS21] suggests that at least the cluster points of the sequence $\{\rho^{*}_{\tau,m}\}$ as $\beta\rightarrow\infty$ coincide with the set of minimizers of the limiting objective ($\beta=\infty$). Another piece of evidence comes from the fact that the annealed dynamics converges to the minimizers of the noiseless objective [Chi22]. We leave for future work the resolution of these issues.

We remark that providing a quantitative bound on the parameter $\tau$ appears to be challenging. The current analysis relies on a dominated convergence argument which does not lead to an explicit convergence rate. Obtaining such a rate requires understanding the trade-off between the terms in the free energy (3.5) for varying $\tau$, and it is also left for future work.

Finally, by combining Theorem 2 with the mean-field analysis in [MMN18], we obtain the desired result on finite-width networks trained via noisy SGD in the low temperature regime.

where the limit in $N,\varepsilon$ is taken along any subsequence $\{(N,\varepsilon=\varepsilon_{N})\}$ with $N/\log\left(N/\varepsilon_{N}\right)\rightarrow\infty$ and $\varepsilon_{N}\log\left(N/\varepsilon_{N}\right)\rightarrow 0$.

In words, Corollary 4.3 means that, at convergence, the estimator implemented by a wide two-layer ReLU network approaches a piecewise linear function, in the regime of vanishingly small noise. In fact, as $\tau,m\to\infty$, the activation function $\sigma^{*}(x,{\bm{\theta}})=a^{\tau,m}(w^{\tau,m}x+b)^{m}_{\tau}$ converges pointwise to the ReLU activation $a(wx+b)_{+}$. We also remark that our result holds for any initialization of the weights of the network, as long as some mild technical conditions are fulfilled (absolute continuity, sub-Gaussian tails, finite second moment and entropy).

Let us clarify some technical aspects of the statement of Corollary 4.3. The result holds for a particular sequence of minimizers, since some of the limits ($t\rightarrow\infty$, $(N,\varepsilon^{-1})\rightarrow\infty$, and $\beta\rightarrow\infty$) are not interchangeable. Furthermore, it appears to be difficult to prove the same statement directly for the noiseless case ($\beta=\infty$). We also point out that the stochasticity of the gradient descent algorithm does not play a role in our analysis, since its impact is seen to be inconsequential by the usual concentration argument [MMN18] when passing to its non-stochastic counterpart.

As concerns the limit in $t$, describing the dependence of the mixing time of the diffusion dynamics (3.1) on the temperature parameter $\beta$ is a cumbersome task. In particular, [GBEK04] show that an exponentially bad dependence could occur if the target function has multiple small risk regions. However, some recent studies show an exponentially fast convergence of the noisy dynamics under some reasonable but particular conditions on the objective landscape [Chi22, NWS22].

As concerns the limit in $(N,\varepsilon)$, the analyses in [MMN18, MMM19] lead to an upper bound on the error term that, with probability at least $1-e^{-z^{2}}$, is given by

where $a\lor b$ denotes the maximum between $a$ and $b$. The exponential dependence of (4.12) in the time $t$ of the dynamics is a common drawback of existing mean-field analyses, and improving it is an open problem which lies beyond the scope of this work. Let us conclude by mentioning that the numerical results presented in Section 7 suggest that, in practical settings, the convergence to the limit occurs rather quickly in the various parameters.

Proof of the Main Results

To do so, in Lemma A.4 we prove that, as $\tau\rightarrow\infty$, the sequence $\rho^{*}_{\tau,m}({\bm{\theta}})$ of minimizers of the free energy $\mathcal{F}^{\tau,m}$ converges pointwise for all ${\bm{\theta}}$ to a minimizer $\rho^{*}_{m}({\bm{\theta}})$ of the free energy $\mathcal{F}^{m}$ with truncated ReLU activation. Then, a dominated convergence argument allows us to obtain (5.1). Next, in Lemma 5.7 we show that, as $\beta\rightarrow\infty$, the absolute value of the integral

can be made arbitrary small for all $x$ except those in the cluster set. The idea is that the absolute value of (5.2) can be bounded by a certain Gaussian integral, and the corresponding covariance matrix is well-defined everywhere except in the cluster set (see Lemmas 5.4 and 5.5). The definition of the cluster set (see (4.2)-(4)) together with the fact that the partition function of $\rho^{*}_{m}$ is uniformly bounded in $m$ (see Lemma 5.2) allows us to show that the cluster set concentrates on at most three points per interval as $\beta\rightarrow\infty$.

In Section 5.3, we show that the predictor evaluated at the Gibbs distribution $\rho^{*}_{\tau,m}$ can be approximated arbitrarily well by an admissible piecewise linear solution (Theorem 2). First, via a Taylor argument, since the curvature vanishes, the estimator can be approximated by a linear function on each interval of $\mathcal{I}\setminus\Omega(m,\beta,\lambda)$. Since the cluster set vanishes concentrating on at most three points per prediction interval, the predictor converges to an admissible piecewise linear solution. However, there is one technical subtlety to consider before reaching this conclusion. Namely, we must consider the possibility that the sequence of predictors experiences unbounded oscillations inside the cluster set, which might ultimately result in a discontinuous limit. Fortunately, this scenario is ruled out because the sequence $\rho^{*}_{\tau,m}$ has uniformly bounded second moments. This fact in conjunction with the structure of the first derivative of the predictor yields the conclusion that the sequence of predictors is equi-Lipschitz, and therefore the limit is well-behaved.

Finally, the proof of Corollary 4.3 follows from similar arguments together with the application of the result established in [MMN18]. More specifically, first, the truncation of the parameter $w$ ensures that, as $t\rightarrow\infty$, the curvature of the predictor evaluated on the solution $\rho_{t}$ of the flow (3.1) converges pointwise in $x$ to the corresponding evaluation on the Gibbs distribution $\rho^{*}_{\tau,m}$. Next, following [MMN18], we couple the weights obtained after $\left\lfloor t/\varepsilon\right\rfloor$ steps of the SGD iteration (3.3) with $N$ i.i.d. particles with distribution $\rho_{t}$, thus obtaining that the curvature of the SGD predictor converges to the curvature of the flow predictor. By using this coupling again, together with the fact that along the trajectory of the flow $M(\rho_{t})<C$ (see [MMN18] or [JKO98]), we obtain a uniform bound on the second moment of the empirical distribution $\hat{\rho}^{N}_{\left\lfloor t/\varepsilon\right\rfloor}$ of the SGD weights. The final result then follows from the same Lipschitz argument described above.

2 Proof of Theorem 1

Let us start with the proof of the vanishing curvature phenomenon. The quantity

is hard to analyze directly due to the presence of the $\tau$-smoothing in the soft-plus activation. However, the structure of the activation $(\cdot)_{\tau}^{m}$ alongside with the pointwise convergence of the minimizers $\rho^{*}_{\tau,m}$ to $\rho^{*}_{m}$ (cf. Lemma A.4) allows us to infer the properties of (5.3) through the analysis of the auxiliary object:

Formally, we show that the approximation result below holds.

Assume that condition A1 holds. Let $\rho^{*}_{\tau,m}$ and $\rho^{*}_{m}$ be the minimizers of the free energy for truncated softplus and ReLU activations, respectively, as defined in (3.9)-(3.10). Then,

where $C$ is a constant independent of $(m,\tau,\beta,\lambda)$.

Recall the definition of the activation $(\cdot)^{\tau}_{m}$ provided in (3.7). We can decompose the integral into two pieces with respect to the domain of truncation and obtain

Let us focus on the first term in the RHS of (5.2). The second derivative has the following form

Thus, the following chain of equalities holds

where in the last step we have performed the change of variables $y=\tau(w^{\tau,m}x+b)$. By Lemma A.4, we have that, as $\tau\to\infty$, $\rho^{*}_{\tau,m}({\bm{\theta}})$ converges to $\rho^{*}_{m}({\bm{\theta}})$ pointwise in ${\bm{\theta}}$. Furthermore, as $\tau\to\infty$, $a^{\tau,m}$ converges to $a^{m}$ for any $a$, and $w^{\tau,m}$ converges to $w^{m}$ for any $w$. Thus, as the Gibbs distributions $\rho^{*}_{\tau,m}({\bm{\theta}})$ and $\rho^{*}_{m}({\bm{\theta}})$ are continuous with respect to ${\bm{\theta}}$, we have that

Furthermore, combining (A.1) and (A.5) from Lemma A.2, we get the following bound

for some constant $C^{\prime}>0$ independent of ${\bm{\theta}}$ and $\tau$. Thus, we have

which is integrable in $(y,a,w)$. Hence, by using the Dominated Convergence theorem and integrating out $y$ using Tonelli’s theorem, we have

Now, by triangle inequality, it remains to show that the absolute value of the second term in the RHS of (5.2) can be upper bounded by $\mathcal{O}\left(\frac{1}{m\lambda}\right)$ as $\tau\rightarrow\infty$. Recall that, by construction,

In addition, we have the following pointwise convergence of the integrand

Furthermore, by using (5.7), we conclude that the integrand can be dominated by an integrable function. Hence, an application of the Dominated Convergence theorem gives that

where the last inequality follows from Lemma A.2, which gives that $M(\rho^{*}_{m})<C^{\prime\prime}/\lambda$ for some $C^{\prime\prime}>0$ that is independent of $(m,\lambda)$. By combining (5.2), (5.8), (5.9) and (5.10), we conclude that (5.5) holds. Finally, by using a standard line of arguments, i.e., Mean Value theorem and Dominated Convergence, the derivative can be pushed inside the integral sign, which finishes the proof. ∎

Next, we study the set on which (5.4) might grow unbounded. In particular, in Lemma 5.3, we provide an upper bound on the measure of the set $\overline{\Omega}_{j}(m,\beta,\lambda)$ defined in (4.5)-(4). To do so, we will first show that the partition function of $\rho^{*}_{m}$ is uniformly bounded in $m$, as stated and proved below.

Consider $\sigma^{*}({\bm{\theta}},x)=a^{\tau,m}(w^{\tau,m}x+b)^{m}_{\tau}$ or $\sigma^{*}({\bm{\theta}},x)=a^{m}(w^{m}x+b)^{m}_{+}$, and let $\rho^{*}_{\sigma^{*}}$ be the Gibbs distribution with activation $\sigma^{*}$. Then, the following upper bound holds for its partition function $Z_{\sigma^{*}}(\beta,\lambda)$:

where $C>0$ is a constant independent of $(m,\tau,\beta,\lambda)$.

Let $R^{\sigma^{*}}_{i}(\rho^{*}_{\sigma^{*}})$ be defined as follows

By substituting the form (3.6) of the Gibbs distribution into the free energy functional (3.5), we have that

Note that, by Fubini’s theorem, we can interchange summation and integration in the first integral, since the activation and the labels are bounded. By using also the definition of $R^{\sigma^{*}}_{i}(\rho_{\sigma^{*}}^{*})$, we have that

where $C>0$ is independent of $(m,\tau,\beta,\lambda)$. From Lemma 10.2 in [MMN18], we obtain that, for any $\rho\in\mathcal{K}$,

where the last inequality follows from non-negativity of $R(\rho)$ and $M(\rho)$. Combining the upper and lower bounds gives

In order to bound the measure of $\overline{\Omega}_{j}(m,\beta,\lambda)$, the idea is to combine the upper bound on the partition function of Lemma 5.2 with a lower bound that diverges in $m$ unless $|\overline{\Omega}_{j}(m,\beta,\lambda)|$ vanishes. In particular, we derive a lower bound with the structure of a Gaussian integral which grows unbounded for a certain set of inputs. This set of inputs corresponds to the scenario when the Gaussian covariance has non-positive eigenvalues, and it can be expressed as the set in which the polynomials $f_{j}$ and $f^{j}$ defined in (4.3) are non-negative. For brevity, we suppress the dependence of $\Omega_{j}$ and $\Omega^{j}$ on $(m,\beta,\lambda)$ in the proofs below.

Assume that condition A1 holds. For $j\in\{0,\ldots,M\}$, let $\Omega^{j}$ and $\Omega_{j}$ be defined as in (4). Then,

where $K_{1},K_{2}>0$ is independent of $(m,\beta,\lambda)$.

We start with the proof for $\Omega^{j}$. For $j=M$, the corresponding polynomial $f^{M}(x)$ is equal to $1+x^{2}$ and therefore $|\Omega^{M}|=0$. Let us now consider the case $j\neq M$, and assume that $\mu(\Omega^{j})>0$. (If that’s not the case, the claim trivially holds.)

Now, we can lower bound the partition function as

Here, the equality in the second line follows from the following observation: if $i\in[j]$ and $(w,b)\in\Omega$, then $w^{m}x_{i}+b\leq 0$ and therefore $(w^{m}x_{i}+b)_{+}^{m}=0$; if $i>j$ and $(w,b)\in\Omega$, then $0<w^{m}x_{i}+b<m^{2}$ ($|x|,|x_{i}|\leq L$, hence $|x_{i}-x|\leq m$, as $L$ is a numerical constant independent of $m$ and $m$ is sufficiently large by assumption A1) and therefore $(w^{m}x_{i}+b)_{+}^{m}=w^{m}x_{i}+b$ for all $(w,b)\in\Omega$. Thus, after the change of variables $(a,w,b)\mapsto(a,w,-w^{m}x)$ and an application of Tonelli’s theorem, the RHS in (5.12) reduces to

Here the coefficients $A^{j}$ and $B^{j}$ are defined as per (4). The term under the exponent can be rewritten as

By definition of $\Omega^{j}$ in conjunction with Sylvester’s criterion, we have that $\Sigma^{-1}$ has a non-positive eigenvalue with corresponding eigenvector

Furthermore, the other eigenvalue with corresponding eigenvector is given by

Note that $v_{-}$ and $v_{+}$ are orthogonal, and consider the following change of variables for the integral

As the matrix $Q$ is unitary, the quantity in (5.13) can be rewritten as

as the determinant of the Jacobian is 1 for any unitary linear transformation. As $\lambda_{-}\leq 0$, this quantity is lower bounded by

Notice that $\|v_{-}\|\geq 1$, $\|v_{+}\|\geq 1$ and $w(\mathbf{z})=z_{1}/\|v_{-}\|_{2}+z_{2}/\|v_{+}\|_{2}.$ Thus, picking $z_{1}\in(0,m/2]$ and $z_{2}\in(0,m/2]$ ensures that $0<w(\mathbf{z})<m$. Furthermore, these conditions on $\mathbf{z}$ do not violate the requirement on $a(\mathbf{z})$, since $|a(\mathbf{z})|\leq|z_{1}|+|z_{2}|\leq m.$ Consequently, as the integrand is non-negative, the integral in (5.14) is lower bounded by

By Lemma A.5, $|R_{i}^{m}(\rho^{*}_{m})|$ is bounded by a constant independent of $(m,\beta,\lambda)$, since $\lambda<C_{3}$ from condition A1. Hence, $\lambda|A^{j}x-B^{j}|$ is also uniformly bounded in $(m,\beta,\lambda)$. This, in particular, implies that

where $K_{1}>0$ is independent of $(m,\beta,\lambda)$. Furthermore, by definition of $\Omega^{j}$, $|B^{j}-A^{j}x|>1$, which implies that $\|v_{+}\|_{2}$ and $\|v_{-}\|_{2}$ are also upper bounded by a constant $K_{2}>0$ independent of $(m,\beta,\lambda)$, and therefore

With this in mind, we can then further lower bound the integral in (5.15) by

In regards to the argument for $\Omega_{j}$, for $j=0$ the result trivially holds, since $f_{0}(x)=1+x^{2}$ and, thus, $|\Omega_{0}|=0$. For $j>0$, the partition function can be lower bounded by

where the set $\Omega$ is defined on non-positive $w$ and $x\in\Omega_{j}$, i.e.,

The rest of the argument remains the same by noting that with the change of variable

which is exactly as in (5.13), but with $x\in\Omega_{j}$ and the polynomial $(B_{j}-A_{j}x)$ in place of $x\in\Omega^{j}$ and the polynomial $(B^{j}-A^{j}x)$. ∎

In order to control the magnitude of (5.4), it is also necessary to understand the behavior of the polynomials defined in (4.3). The worst case scenario, in terms of presenting a challenge to bounding the curvature, corresponds to $f^{j}$ or $f_{j}$ being arbitrarily close to zero on the whole area outside of cluster set. In fact, this would imply that the Gaussian-like integral arising in the computation of (5.4) has arbitrary small eigenvalues. More specifically, our plan is to exploit the following bound for $x\in I_{j}\setminus\overline{\Omega}_{j}(m,\beta,\lambda)$:

Now, the RHS of (5.18) diverges (and, therefore, the bound is useless), if either of the polynomials is arbitrarily close to zero outside of the cluster set. Fortunately, we are able to prove that this cannot happen: in Lemma 5.5 we show that $f^{j}(x)$ and $f_{j}(x)$ can be small only when $x$ approaches the cluster set, i.e.,

where $C^{j}(x),C_{j}(x)$ are defined in (4.8) and, because of the condition on their coefficients $\{K_{i}\}_{i=1}^{4}$, they cannot be arbitrarily close to in any interval $I_{j}$.

where we suppress the dependence on $(a,b)$, i.e., $P_{2}(x;a,b)=P_{2}(x)$, for more compact notation. In addition, let $\Omega_{+}$ be the subset of $I$ on which $P_{2}$ is strictly positive, i.e.,

For a fixed small constant $C_{\Omega}>0$, define the set of admissible coefficients as follows

Given $(a,b)\in\mathcal{U}$ and $x\in\Omega_{+}$, we define the critical point $x_{c}$ of the polynomial $P_{2}$ associated with $x$ and $\Omega_{+}$ in the same fashion as in Definition 4.1, after replacing $f^{j}(\cdot)$ with $P_{2}(\cdot)$ and $I_{j}\setminus\Omega^{j}(m,\beta,\lambda)$ with $\Omega_{+}$. Notice that, since $\Omega_{+}$ has strictly positive Lebesgue measure for $(a,b)\in\mathcal{U}$, the critical point is well-defined and, in particular, $x_{c}\in I$ always holds.

Fix some $C_{\Omega}$ such that $\mathcal{U}$, as defined in (5.20), is of positive measure. Pick some interval $(a,b)\in\mathcal{U}$. Let $x\in\Omega_{+}$ and $x_{c}$ be the critical point associated to $x$. Then, the following holds

where $\alpha_{0},\alpha_{1},\alpha_{2}\geq 0$ and at least one of them is lower bounded by a strictly positive constant depending on $C_{\Omega}$ but independent of the choice of $(a,b)\in\mathcal{U}$.

We defer the proof of Lemma 5.4 to Appendix A.3. Recall the definition of the polynomial $f^{j}(x)$ given in (4.3), and notice that expression can be rearranged such that $f^{j}(x)$ is in the form of (5.19), namely

In this view, the following result follows from Lemma 5.4.

Assume that $(A^{j},B^{j})\in\mathcal{U}$, i.e., $|I_{j}\setminus\Omega^{j}|$ is lower bounded by a positive constant. Given $x\in I_{j}\setminus\Omega^{j}$, let $x_{c}$ be the critical point associated to $x$. Then, we have that

where $\gamma_{1},\gamma_{2}>0$ and either $\gamma_{1}>\varepsilon$ or $\gamma_{2}>\varepsilon$ for some $\varepsilon>0$ that is independent of $(A^{j},B^{j})$ but depending on $C_{\Omega}$ as appearing in the definition of $\mathcal{U}$.

Note that $I_{j}\setminus\Omega^{j}$ is the set in which $f^{j}$ is strictly positive. Hence, since $|I_{j}\setminus\Omega^{j}|$ is lower bounded by a positive constant independent of $A^{j},B^{j}$, we can apply Lemma 5.4 to get

where $\alpha_{0},\alpha_{1},\alpha_{2}\geq 0$ and at least one of them is lower bounded by a strictly positive constant independent of $(A^{j},B^{j})$. Thus, since each term of the RHS above is non-negative, we get

where $i=\arg\max_{j\in\{1,2\}}\alpha_{j}$. Furthermore, as $|x-x_{c}|\leq|I_{j}|$, we have

Now, either $\alpha_{i}$ or $\alpha_{0}$ as well as $1/|I_{j}|$ are lower bounded by strictly positive constants independent of $(A^{j},B^{j})$. Thus, taking $\gamma_{1}=\alpha_{i}/|I_{j}|^{2-i}$ and $\gamma_{2}=\alpha_{0}$ concludes the proof. ∎

Let us point out that, although $\varepsilon$ does not depend on the values of $(A^{j},B^{j})\in\mathcal{U}$, the position of a critical point $x_{c}$ depends on $(A^{j},B^{j})$.

In a similar fashion, we define $\bar{\mathcal{U}}$ to be the set of admissible $(A_{j},B_{j})$ as in (5.20), and given $x\in I_{j}\setminus\Omega_{j}$, we let $\bar{x}_{c}$ be the critical point associated to $x$ and $\Omega_{j}$. Then, a result analogous to Lemma 5.5 holds for $f_{j}(x)$:

where $\gamma_{3},\gamma_{4}>0$ and either $\gamma_{3}>\varepsilon$ or $\gamma_{4}>\varepsilon$ for some $\varepsilon>0$ that is independent of the choice of $(A_{j},B_{j})\in\bar{\mathcal{U}}$.

The last ingredient for the proof of the vanishing curvature phenomenon is the control of the decay of the partition function $Z_{m}(\beta,\lambda)$ as $\beta\rightarrow 0$.

for some $C>0$ that is independent of $(m,\beta,\lambda)$.

The proof of Lemma 5.6 is deferred to Appendix A.2. At this point, we are ready to provide an upper bound on the magnitude of (5.4).

Assume that condition A1 holds. Furthermore, assume that $m>e^{\beta K_{2}}$, where $K_{2}$ is given in (5.11). Fix $j\in\{0,\ldots,M\}$. Then, for any $x\in I_{j}\setminus\left(\Omega^{j}\cup\Omega_{j}\right)$,

where $K>0$ is independent of $(m,\beta,\lambda)$, $\bar{C}^{j}(x):=\min\left\{C_{j}(x),C^{j}(x),1\right\}$, and $C^{j}(x),C_{j}(x)$ are given by (5.22) and (5.23), respectively.

Note that the following upper bound holds

Let us now decompose the integral $I(x)$ depending on the sign of $w$, i.e.,

and, recalling the form of $\rho^{*}_{m}(a,w,-w^{m}x)$ from (3.10), the corresponding potentials are given by

By recalling from (4) the definitions of $A^{j},A_{j},B^{j}$ and $B_{j}$, we obtain the following upper bounds.

Let us analyze the RHS of (5.24). This term can be rewritten as

where the first inequality follows from Lemma 5.3, and the second inequality uses that $m>e^{\beta K_{2}}$. Therefore, for sufficiently large $\beta$, $|\Omega^{j}|$ is smaller than $|I_{j}|/2$, and therefore $|I_{j}\setminus\Omega^{j}|$ is lower bounded by $|I_{j}|/2$. At this point, we can apply Lemma 5.5 which gives that $1+x^{2}-(A^{j}x-B^{j})^{2}\geq C^{j}(x)\geq\bar{C}^{j}(x):=\min\left\{C^{j}(x),C_{j}(x),1\right\}$. Thus, (5.26) is upper bounded by

where $A\sim\mathcal{N}(|B^{j}-A^{j}x|w^{m},(\beta\lambda)^{-1})$. Furthermore, the following chain of inequalities hold:

where the second passage follows from Jensen’s inequality. By using (5.28), the RHS of (5.27) is upper bounded by

Applying Lemma 5.5 again to obtain $(A^{j}x-B^{j})^{2}\leq 1+x^{2}-\bar{C}^{j}(x)\leq 1+x^{2}$ and noting by definition that $(w^{m})^{2}\leq w^{2}$, we now upper bound this last term by

where in the second line we use that $1-\bar{C}^{j}(x)\geq 0$ and again that $(w^{m})^{2}\leq w^{2}$, and in the third line we use that $\sqrt{u+v}\leq\sqrt{u}+\sqrt{v}$.

Finally, computing explicitly the last integral gives the following upper bound on the RHS of (5.24) and consequently on $I^{j}(x)$:

By following the similar passages, we obtain the same upper bound for $I_{j}(x)$. By using the lower bound on the partition function shown in Lemma 5.6, we conclude that

where $K>0$ is independent of $(m,\beta,\lambda)$, and the proof is complete. ∎

The proof of Theorem 1 is an immediate consequence of the results presented so far.

The proof of (4.9) follows from Lemmas 5.1 and 5.7, and the proof of (4.11) follows from Lemma 5.3. ∎

3 Proof of Theorem 2

To summarize, at this point we have shown that as $\beta\to\infty$ the second derivative of the predictor vanishes outside the cluster set, and that the size of the cluster set shrinks to concentrate on at most 3 points per prediction interval. With these results in mind, we are ready to provide the proof for Theorem 2.

The predictor evaluated at the Gibbs distribution is given by

where $n=(\tau,m,\beta,\lambda)$ denotes the aggregated index and we suppress the dependence on $(\beta,\lambda)$ in $\rho^{*}_{\tau,m}$ for convenience. By Lemma A.6, there exists $\tau(m,\beta,\lambda)$ such that, for any $\tau>\tau(m,\beta,\lambda)$,

for some $C>0$ independent of $(\tau,m,\beta,\lambda)$. We start by showing that the family of predictors $\{y_{n}\}$ is equi-Lipschitz for $\infty>\tau>\tau(m,\beta,\lambda)$. First, note that

since the derivative can be pushed inside by the same line of arguments as given in the proof of Lemma 5.1. Next, we have that, by construction of the activation, the following holds

Using the Cauchy-Schwartz inequality and (5.30), we obtain that

By combining (5.31) and (5.32), we have shown that the family $\{y_{n}\}$ for $\tau>\tau(m,\beta,\lambda)$ is equi-Lipschitz, as the derivatives are uniformly bounded. By using a similar argument, we can show that the same result holds for the predictor itself, i.e., for all $x\in\bigcup_{j=0}^{M}I_{j}$, $y_{n}(x)$ is uniformly bounded.

Note that Theorem 1 considers the curvature of points outside the cluster set, and it gives an upper bound which diverges when $\bar{C}^{j}(x)$ approaches for some $j\in[M]$. Thus, our next step is to develop the analytical machinery to make this scenario impossible. Let us recall Definitions (4.8) and (4.10). Then, by Lemma 5.5, we have that

where $\gamma_{1},\gamma_{2},\gamma_{3},\gamma_{4}>0$ and $\min\{\max\{\gamma_{1},\gamma_{2}\},\max\{\gamma_{3},\gamma_{4}\}\}>\varepsilon$, for some $\varepsilon>0$ that is independent of $(m,\beta,\lambda)$. Let us focus on the term $\gamma_{1}(x-x_{c})^{2}+\gamma_{2}$. If $\gamma_{2}=0$ or it approaches (as $m,\beta\to\infty$), then we extend $\Omega^{j}(m,\beta,\lambda)$ as

where $\{A_{i}^{j}\}_{i=1}^{3}$ denote such (possibly empty) open intervals. Furthermore, $I_{j}\setminus\bar{\Omega}^{j}_{\rm ext}$ is the union of at most three disjoint closed intervals, i.e.,

where $\{B_{i}^{j}\}_{i=1}^{3}$ denote such (possibly empty) closed intervals.

where we use that $|x-x_{c}|\geq\delta$ by construction of the extended set $\bar{\Omega}^{j}_{\rm ext}$. Let us define

Then, by picking a sufficiently small $\delta$, (5.33) implies that, as $m\lambda\rightarrow\infty$ and $\beta\lambda^{7/4}\rightarrow\infty$, for all $x\in B_{i}^{j}$,

Let us now consider the open intervals $\{A_{i}^{j}\}_{i=1}^{3}$. For any $x\in A_{i}^{j}$, let

where the second inequality is due to the fact that the families $\{y_{n}\}$ and $\{f_{n}^{i}\}$ are equi-Lipschitz. From (5.34) the second term in the RHS in (5.3) vanishes. As for the first term, by construction of the extension, together with the result of Lemma 5.3, we have that

for some $K_{2}>0$ independent of $(m,\beta,\lambda)$. Thus, by picking a sufficiently small $\delta$ and $m>e^{\beta K_{2}}$, we conclude that the first term in the RHS in (5.3) also vanishes.

So far, we have showed that, both inside and outside of the extension of the cluster set, the predictor $y_{n}$ is well approximated by linear functions. It remains to prove that the linear pieces connect, i.e., there exists $\hat{x}\in\bar{\Omega}^{j}_{\rm ext}$ such that, for two neighboring linearities $f_{n}^{i}$ and $f_{n}^{i+1}$ (possibly belonging to different intervals), the following holds

This claim follows from Lipschitz arguments similar to those presented above, and the proof is complete. ∎

4 Proof of Corollary 4.3

At this point, we have proved a result about the structure of the predictor coming from the minimizer of the free energy (3.5). By using the mean-field analysis in [MMN18], we finally show that this structural result holds for the predictor obtained from a wide two-layer ReLU network.

First, we show that, as $t\to\infty$, the second derivative of the predictor evaluated on the solution $\rho_{t}$ of the flow (3.1) converges to the same quantity evaluated on the Gibbs minimizer $\rho^{*}_{\tau,m}$. To do so, we decompose the integral involving $\rho_{t}$ as in Lemma 5.1 (cf. (5.2)):

Next, we show that a technical condition bounding the free energy at initialization appearing in the statement of Theorem 4 in [MMN18] is satisfied under the assumption $M(\rho_{0})<\infty$ and $H(\rho_{0})>-\infty$. Recalling the sandwich bound for the truncated soft-plus activation (3.8) and the fact that that $\tau\geq 1$ by condition A1, an application of Cauchy-Schwarz inequality gives

where $C,C^{\prime}>0$ are some numerical constants independent of $(\tau,m)$. This readily implies that

since $\lambda$ and $\beta^{-1}$ are upper-bounded by assumption A1.

Now we can apply Theorem 4 in [MMN18] to conclude that, as $t\rightarrow\infty$,

Thus, as the terms inside the integrals in (5.4) are all bounded for fixed $(\tau,m,\beta,\lambda)$, by definition of weak convergence, we get that, as $t\rightarrow\infty$,

Consequently, since the derivative operator can be pushed inside by the same arguments as in Lemma 5.1, we have that, as $t\rightarrow\infty$, the following pointwise convergence holds

Next, we show that the second derivative of the predictor obtained from the two-layer ReLU network also converges to the same limit. Recall that $\sigma^{*}(x,{\bm{\theta}})=a^{\tau,m}(w^{\tau,m}x+b)^{m}_{\tau}$. Then, by Theorem 3 in [MMN18], we have that, almost surely, as $N\rightarrow\infty,\,\,\varepsilon_{N}\rightarrow 0$

along any sequence $\{\varepsilon_{N}\}$ such that $\varepsilon_{N}\log(N/\varepsilon_{N})\rightarrow 0$ and $N/\log(N/\varepsilon_{N})\rightarrow\infty$. By combining (5.37) and (5.38), we obtain that the desired convergence result holds for the LHS of (5.37).

Another application of Theorem 3 of [MMN18], together with the fact that the second moment of the flow solution $\rho_{t}$ is uniformly bounded along the sequence $t\rightarrow\infty$ (cf. Lemma 10.2 in [MMN18], following Proposition 4.1 in [JKO98]), gives that the gradients

are almost surely uniformly bounded. This fact, in turn, implies that the corresponding predictor is almost surely equi-Lipschitz. In a similar fashion, we also have that the predictor itself is almost surely uniformly bounded in absolute value.

At this point, the desired result follows from the same line of arguments as in the proof of Theorem 2. ∎

Knots Inside the Interval

In this section, we provide an explicit example of a 2-point dataset such that the SGD solution exhibits a change of tangent (or “knot”) inside the training interval. To do so, we will show that neural networks implementing a linear function without knots on the prediction interval cannot minimize the free energy (3.5). To simplify the analysis, throughout the section we omit the limits in $(\tau,m)$, i.e., we consider directly ReLU activations (this corresponds to taking $\tau=m=\infty$). Similar arguments apply to the case of sufficiently large parameters $\tau$ and $m$.

We start with the case of noiseless SGD training, i.e., $\beta=+\infty$. Here, the free energy has no entropy penalty and it can be expressed as

We consider the following dataset which consists of two points:

Let $f^{*}(x)$ be the piecewise linear function that interpolates the training data $\{(-\bar{x},\bar{y}),(\bar{x},\bar{y})\}$ and passes through the point $(0,0)$, where it exhibits a knot (see the orange curve in Figure 6a). Note that

and $\delta_{(a_{0},w_{0},b_{0})}$ denotes the Dirac delta function centered at $(a_{0},w_{0},b_{0})$. Note that $R(\rho^{*})=0$ and $M(\rho^{*})=\frac{2}{5}$. Thus, the free energy is given by

Let $f(x)$ be a linear function on the interval $[-\bar{x},\bar{x}]$ such that $f(-\bar{x})=\bar{y}+\varepsilon_{1}$ and $f(\bar{x})=\bar{y}+\varepsilon_{2}$ (see the blue dashed line in Figure 6a), and let $\rho$ be the corresponding distribution of the parameters, i.e.,

In the rest of this section, we will show that, for all $\lambda\leq 1$,

In words, the minimizer of the free energy cannot be a linear function on the interval $[-\bar{x},\bar{x}]$. As $f$ is linear, we have that

First, we consider the case $f(0)=0$. From (6.7), we have that $\varepsilon_{1}+\varepsilon_{2}=-2\bar{y}$. Hence,

By combining (6.8) and (6.4), we conclude that (6.6) holds for all $\lambda\leq 1$ (under the additional restriction $f(0)=0$).

Next, we consider the case $f(0)\neq 0$. By using (6.7) and applying Cauchy-Schwarz inequality, we have that

With this in mind, we can lower bound the regularization term as

where the last inequality follows from the fact that $g(t)=(f(0))^{2}/t+t$ is minimized over $t\geq 0$ by taking $t=|f(0)|$. Therefore, we have that

Note that, for a fixed value of the sum $\varepsilon_{1}+\varepsilon_{2}$, the quantity $\varepsilon_{1}^{2}+\varepsilon_{2}^{2}$ is minimized when $\varepsilon_{1}=\varepsilon_{2}$. Thus, by recalling that $\bar{y}=2$, we have

One can readily verify that, for any $\lambda\leq 2$, the minimizer is given by $\varepsilon^{*}=-\lambda$. Thus,

where the first inequality uses (6.9) and that the minimizer is $\varepsilon^{*}=-\lambda$, and the next two inequalities use that $\lambda\geq 1$. Merging two cases regarding $f(0)$, we conclude that (6.6) holds, as desired.

2 Low Temperature Regime

We now focus on the case of noisy SGD with temperature $\beta^{-1}$. Here, the free energy can be expressed as

We consider the two-point dataset (6.2) and we recall that $f^{*}(x)$ has a knot inside the training interval. In this section we will show that the following two results hold for all $\lambda\leq 1$:

There exists a sequence of distributions $\{\rho^{*}_{\beta}\}_{\beta}$ such that, for any $x\in[-\bar{x},\bar{x}]$,

Let $\rho$ be a distribution such that the function $f(x)$ given by (6.5) is linear in the interval $[-\bar{x},\bar{x}]$. Pick a sequence of distributions $\{\rho_{\beta}\}_{\beta}$ such that $\rho_{\beta}\rightharpoonup\rho$ and for any $x\in[-\bar{x},\bar{x}]$,

Combining these two results gives that, for sufficiently large $\beta$, the minimizer of the free energy (6.11) cannot yield a linear estimator on the interval between the two data points. In Figure 6b, we represent the function obtained by training via SGD a two-layer ReLU network with 500 neurons on the dataset (6.2). Clearly, the blue curve approaches the piecewise linear function $f^{*}(x)$, which contains a knot inside the interval $$. The plot represented in the Figure corresponds to the case with no regularization ($\lambda=0$), but similar results are obtained for small (but non-zero) regularization.

where we have used that the support of $\rho^{*}$ lies inside the ball $B_{r}$. The first term in the RHS of (6.16) vanishes as $\beta\rightarrow\infty$ by weak convergence, since the function $a(wx+b)_{+}$ is bounded inside $B_{r}$. For the second term, we have that, for any $x\in[-\bar{x},\bar{x}]$,

where $C>0$ is a constant independent of $(\beta,r)$. Since the sequence $\{\|X_{\beta}\|_{2}^{2}\}_{\beta}$ is uniformly integrable, we can make the RHS arbitrary small by picking a sufficiently large $r$ (uniformly for all $\beta$). As a result, (6.12) readily follows.

Note that (6.12) immediately implies that, as $\beta\rightarrow\infty$, $R(\rho^{*}_{\beta})\rightarrow R(\rho^{*})=0$. Furthermore, with similar arguments we obtain that, as $\beta\to\infty$, $M(\rho^{*}_{\beta})\to M(\rho^{*})$. By convexity of the differential entropy, we have that $H(\frac{1}{2}\rho_{1}+\frac{1}{2}\rho_{2})\geq\frac{1}{2}H(\rho_{1})+\frac{1}{2}H(\rho_{2})$. Hence, $H(\rho_{\beta}^{*})\geq C\log(2\pi e/\beta)$, where $C>0$ is independent of $\beta$. By combining these bounds on $R(\rho^{*}_{\beta})$, $M(\rho^{*}_{\beta})$ and $H(\rho^{*}_{\beta})$, we conclude that

which, combined with (6.4), completes the proof of (6.13).

Proof of (ii).

From (6.14), we obtain that $\lim_{\beta\to\infty}R(\rho_{\beta})=R(\rho)$. As the second moment is lower-semicontinuous and bounded from below, we have that $\liminf_{\beta\to\infty}M(\rho_{\beta})\geq M(\rho)$. Furthermore, Lemma 10.2 in [MMN18] implies that

By replicating the argument leading to (6.10) (but now with regularization coefficient $\lambda/2$ instead of $\lambda$), we obtain that the RHS of (6.17) can be lower bounded as

for all $\lambda\leq 1$. Then, the desired result follows from (6.17) and (6.18).

Numerical Simulations

We consider training the two-layer neural network (3.1) with $N$ neurons and ReLU activation functions, i.e., $\sigma^{*}(x,{\bm{\theta}})=a(wx+b)_{+}$, with ${\bm{\theta}}=(a,w,b)$. We run the SGD iteration (3.3) (no momentum or weight decay, batch size equal to $1$), and we plot the resulting predictor once the algorithm has converged. The results for two different unidimensional datasets are reported in Figures 7 and 8. In these experiments, we set $N=500$ and we remark that the plots for wider networks ($N\in\{1000,2000,5000\}$) look identical. We also point out that the shape of the predictor does not change for different runs of the SGD algorithm (with different initializations, and order of the training samples). This is in agreement with the mean-field predictions when $\beta<\infty$, $\lambda>0$ and the variance of the initialization does not depend on $N$. The same setup is employed to obtain the numerical results of Figure 1 and 6b, discussed in Section 1 and 6, respectively.