Particle Dual Averaging: Optimization of Mean Field Neural Networks with Global Convergence Rate Analysis

Introduction

Gradient-based optimization can achieve vanishing training error on neural networks, despite the apparent non-convex landscape. Among various works that explains the global convergence, one common ingredient is to utilize overparameterization to translate the training dynamics into function spaces, and then exploit the convexity of the loss function with respect to the function. Such endeavors usually consider models in one of the two categories: the mean field regime or the kernel regime.

On one hand, analysis in the kernel (lazy) regime connects gradient descent on wide neural network to kernel regression with respect to the neural tangent kernel (Jacot et al.,, 2018), which leads to global convergence at linear rate (Du et al.,, 2019; Allen-Zhu et al.,, 2019; Zou et al.,, 2020). However, key to the analysis is the linearization of the training dynamics, which requires appropriate scaling of the model such that distance traveled by the parameters vanishes (Chizat and Bach, 2018a, ). Such regime thus fails to explain the feature learning of neural networks (Yang and Hu,, 2020), which is believed to be an important advantage of deep learning; indeed, it has been shown that deep learning can outperform kernel models due to this adaptivity (Suzuki,, 2019; Ghorbani et al., 2019a, ).

In contrast, the mean field regime describes the gradient descent dynamics as Wasserstein gradient flow in the probability space (Nitanda and Suzuki,, 2017; Mei et al.,, 2018; Chizat and Bach, 2018b, ), which captures the potentially nonlinear evolution of parameters travelling beyond the kernel regime. While the mean field limit is appealing due to the presence of “feature learning”, its characterization is more challenging and quantitative analysis is largely lacking. Recent works established convergence rate in continuous time under modified dynamics (Rotskoff et al.,, 2019), strong assumptions on the target function (Javanmard et al.,, 2019), or regularized objective (Hu et al.,, 2019), but such result can be fragile in the discrete-time or finite-particle setting — in fact, the discretization error often scales exponentially with the time horizon or dimensionality, which limits the applicability of the theory. Hence, an important research problem that we aim to address is

Can we develop optimization algorithms for neural networks in the mean field regime with more accurate quantitative guarantees the kernel regime enjoys?

We address this question by introducing the particle dual averaging (PDA) method, which globally optimizes an entropic regularized nonlinear functional. For two-layer mean field network which is an important application, we establish polynomial runtime guarantee for the discrete-time algorithm; to our knowledge this is the first quantitative global convergence result under similar settings.

We propose the PDA algorithm, which draws inspiration from the dual averaging method originally developed for finite-dimensional convex optimization (Nesterov,, 2005, 2009; Xiao,, 2009). We iteratively optimize a probability distribution in the form of a Boltzmann distribution, samples from which can be obtained from the Langevin algorithm (see Figure 1.1). The resulting algorithm has comparable per-iteration cost as gradient descent and can be efficiently implemented.

2 Related Literature

The key observation for the mean field analysis is that when the number of neurons becomes large, the evolution of parameters is well-described by a nonlinear partial differential equation (PDE), which can be viewed as solving an infinite-dimensional convex problem (Bengio et al.,, 2005; Bach,, 2017). Global convergence can be derived by studying the limiting PDE (Mei et al.,, 2018; Chizat and Bach, 2018b, ; Rotskoff and Vanden-Eijnden,, 2018; Sirignano and Spiliopoulos,, 2020), yet quantitative convergence rate generally requires additional assumptions.

Our setting is most related to Hu et al., (2019), who studied the minimization of a nonlinear functional with KL regularization on the probability space, and showed linear convergence (in continuous time) of a particle dynamics named mean field Langevin dynamics when the regularization is sufficiently strong. Chen et al., (2020) also considered optimizing a KL-regularized objective in the infinite-width and continuous-time limit, and derived NTK-like convergence guarantee under certain parameter scaling. Compared to these prior works, we directly handle the discrete time update in the mean-field regime, and our analysis covers a wider range of regularization parameters and loss functions.

Langevin dynamics can be viewed as optimization in the space of probability measures (Jordan and Kinderlehrer,, 1996; Jordan et al.,, 1998); this perspective has been explored in Wibisono, (2018); Durmus et al., (2019). It is known that the continuous-time Langevin diffusion converges exponentially fast to target distributions satisfying certain growth conditions (Roberts and Tweedie,, 1996; Mattingly et al.,, 2002). The discretized Langevin algorithm has a sublinear convergence rate that depends on the numerical scheme (Li et al.,, 2019) and has been studied under various metrics (Dalalyan,, 2014; Durmus and Moulines,, 2017; Cheng and Bartlett,, 2017).

The Langevin algorithm can also optimize certain non-convex objectives (Raginsky et al.,, 2017; Xu et al.,, 2018; Erdogdu et al.,, 2018), in which one finite-dimensional “particle” can attain approximate global convergence due to concentration of Boltzmann distribution around the true minimizer. However, such result often depends on the spectral gap that grows exponentially in dimensionality, which renders the analysis ineffective for neural net optimization in the high-dimensional parameter space.

Very recently, convergence of Hamiltonian Monte Carlo in learning certain mean field models has been analyzed in Bou-Rabee and Schuh, (2020); Bou-Rabee and Eberle, (2021). Compared to these concurrent results, our formulation covers a more general class of potentials, and in the context of two-layer neural network, we provide optimization guarantees for a wider range of loss functions.

3 Notations

Problem Setting

A typical example in the literature of the above formulation is a two-layer neural network.

The purpose of $\sigma_{2}$ in the last layer is to ensure the boundedness of output (e.g., see Assumption 2 in Mei et al., (2018)); this nonlinearity can also be removed if parameters of output layer are fixed. In addition, although we mainly focus on the optimization of two-layer neural network, our proposed method can also be applied to ensemble $h_{\Theta}$ of deep neural networks $h_{\theta_{r}}$.

2 Regularized Empirical Risk Minimization

where $R_{\lambda_{1},\lambda_{2}}$ is a regularization term composed of the weighted sum of the second-order moment and negative entropy with regularization parameters $\lambda_{1}$, $\lambda_{2}$:

Note that this regularization is the KL divergence of $q$ from a Gaussian distribution. In our setting, such regularization ensures that the Gibbs distributions $q_{*}^{(t)}$ specified in Section 3 are well defined.

While our primary focus is the optimization of the objective (3), we can also derive a generalization error bound for the empirical risk minimizer of order of $O(n^{-1/2})$ for both the regression and binary classification settings, following Chen et al., (2020). We defer the details to Appendix D.

3 The Langevin Algorithm

Under appropriate conditions on $f$, it is known that $\theta^{(t)}$ converges to a stationary distribution proportional to $\exp(-f(\cdot))$ in terms of KL divergence at a linear rate (e.g., Vempala and Wibisono, (2019)) up to $O(\eta)$-error, where we hide additional factors in the big-$O$ notation.

Hence we may interpret the Langevin algorithm as approximately solving an entropic regularized linear functional (i.e., free energy functional) on the probability space. This connection between sampling and optimization (see Dalalyan, (2017); Wibisono, (2018); Durmus et al., (2019)) enables us to employ the Langevin algorithm to obtain (samples from) the closed-form Boltzmann distribution which is the minimizer of (6); for example, many Bayesian inference problems fall into this category.

Proposed Method

We now propose the particle dual averaging method to approximately solve the problem (3) by optimizing a two-layer neural network in the mean field regime; we also introduce the mean field limit of the proposed method to explain the algorithmic intuition and develop the convergence analysis.

Our proposed particle dual averaging method (Algorithm 1) is an optimization algorithm on the space of probability measures. The algorithm consists of an inner loop and outer loop; we run Langevin algorithm in inner loop to approximate a Gibbs distribution, which is optimized in the outer loop so that it converges to the optimal distribution $q_{*}$. This outer loop update is designed to extend the classical dual averaging scheme (Nesterov,, 2005, 2009; Xiao,, 2009) to infinite dimensional optimization problems (described in Section 3.2). Below we provide a more detailed explanation.

2 Mean Field View of PDA

In this subsection we discuss the mean field limit of PDA and explain its algorithmic intuition. Note that the inner loop of Algorithm 1 is the Langevin algorithm with $M$ particles, which optimizes the potential function given by the weighted sum:

Due to the rapid convergence of Langevin algorithm outlined in Subsection 2.3, the particles $\theta_{r}^{(k+1)}$ $(r\in\{1,\ldots,M\})$ can be regarded as (approximate) samples from the Boltzmann distribution: $\exp\left(-\overline{g}^{(t)}\right)$. Hence, the inner loop of PDA returns an $M$-particle approximation of some stationary distribution, which is then modified in the outer loop. Importantly, the update on the stationary distribution is designed so that the algorithm converges to the optimal solution of the problem (3).

We now introduce the mean field limit of PDA, i.e., taking the number of particles $M\to\infty$ and directly optimizing the problem (3) over $q$. We refer to this mean field limit simply as the dual averaging (DA) algorithm. The dual averaging method was originally developed for the convex optimization in finite-dimensional spaces (Nesterov,, 2005, 2009; Xiao,, 2009), and here we adapt it to optimization on the probability space. The detail of the DA algorithm is described in Algorithm 2.

Algorithm 2 iteratively updates the density function $q^{(t+1)}_{*}\in\mathcal{P}_{2}$ which is a solution to the objective:

To sum up, we may interpret the DA method as approximating the optimal distribution $q_{*}$ by iteratively optimizing $q_{*}^{(t)}$, which takes the form of a Boltzmann distribution. In the inner loop of the PDA algorithm, we obtain $M$ (approximate) samples from $q_{*}^{(t)}$ via the Langevin algorithm. In other words, PDA can be viewed as a finite-particle approximation of DA – indeed, the stationary distributions obtained in PDA converges to $q_{*}^{(t+1)}$ by taking $M\rightarrow\infty$. In the following section, we present the convergence rate of the DA method, and also take into account the iteration complexity of the Langevin algorithm; we defer the finite-particle approximation error analysis to Appendix C.

Convergence Analysis

We now provide quantitative global convergence guarantee for our proposed method in discrete time. We first derive the outer loop complexity, assuming approximate optimality of the inner loop iterates, which we then verify in the inner loop analysis. The total complexity is then simply obtained by combining the outer- and inner-loop runtime.

We first analyze the convergence rate of the dual averaging (DA) method (Algorithm 2). Our analysis will be made under the following assumptions.

(A2) $|h(\theta,x)|\leq 1$ and $h(\theta,x)$ is smooth with respect to $\theta$ for $x\in\mathcal{X}$.

(A2) is satisfied by smooth activation functions such as sigmoid and tanh. Many loss functions including the squared loss and logistic loss satisfy (A1) under the boundedness assumptions $\mathcal{Y}\subset$ and $|h_{\theta}(x)|\leq 1$. Note that constants in (A1) and (A2) are defined for simplicity and can be relaxed to any value. (A3) specifies the precision of approximate solutions of sub-problems (7) to guarantee the global convergence of Algorithm 2, which we verify in our inner loop analysis.

We first introduce the following quantity for $q\in\mathcal{P}_{2}$,

Observe that the expression consists of the negative entropy minus its lower bound for $q^{(t)}_{*}$ under Assumption (A1), (A2); in other words $e(q^{(t)}_{*})\geq 0$. We have the following convergence rate of DAIn Appendix B we introduce a more general version of Theorem 1 that allows for inexact $h_{q^{(t)}}(x)$, which simplifies the analysis of finite-particle discretization presented in Appendix C..

Under Assumptions (A1), (A2), and (A3), for arbitrary $q_{*}\in\mathcal{P}_{2}$, iterates of the DA method (Algorithm 2) satisfies

On the other hand, due to the exponential dependence on $\lambda_{2}^{-1}$, our convergence rate is not informative under weak regularization $\lambda_{2}\to 0$. Such dependence follows from the classical LSI perturbation lemma (Holley and Stroock,, 1987), which is likely unavoidable for Langevin-based methods in the most general setting (Menz and Schlichting,, 2014), unless additional assumptions are imposed (e.g., a student-teacher setup); we intend to further investigate these conditions in future work.

2 Inner Loop Complexity

In order to derive the total complexity (i.e., taking both the outer loop and inner loop into account) towards a required accuracy, we also need to estimate the iteration complexity of Langevin algorithm. We utilize the following convergence result under the log-Sobolev inequality (Definition A):

(A4) $\partial_{\theta}h(\cdot,x)$ is $1$-Lipschitz continuous: $\|\partial_{\theta}h(\theta,x)-\partial_{\theta}h(\theta^{\prime},x)\|_{2}\leq\|\theta-\theta^{\prime}\|_{2}$, $\forall x\in\mathcal{X}$, $\theta,\theta^{\prime}\in\Omega$.​​

(A4) is parallel to (Mei et al.,, 2018, Assumption A3), and is satisfied by two-layer neural network in Example 1 when the output or input layer is fixed and the input space $\mathcal{X}$ is compact. We remark that this assumption can be relaxed to Hölder continuity of $\partial_{\theta}h(\cdot,x)$ via the recent result of Erdogdu and Hosseinzadeh, (2020), which allows us to extend Theorem 1 to general $L_{p}$-norm regularizer for $p>1$. For now we work with (A4) for simplicity of the presentation and proof.

Set $\delta_{t+1}$ to be the desired accuracy of an approximate solution $q^{(t+1)}$ specified in (A3): $\delta_{t+1}=1/(t+1)^{2}$, we have the following guarantee for the inner loop.

We comment that for the inner loop we utilized the overdamped Langevin algorithm, since it is the most standard and commonly used sampling method for the objective (7). Our analysis can easily incorporate other inner loop updates such as the underdamped Langevin algorithm (Cheng et al.,, 2018; Eberle et al.,, 2019) or the Metropolis-adjusted Langevin algorithm (Roberts and Tweedie,, 1996; Dwivedi et al.,, 2018), which may improve the iteration complexity.

3 Total Complexity

Experiments

We employ our proposed algorithm in both synthetic student-teacher settings (see Figure 1(a)(b)) and real-world dataset (see Figure 1(c)). For the student-teacher setup, the labels are generated as $y_{i}=f_{*}(x_{i})+\varepsilon_{i}$, where $f_{*}$ is the teacher model (target function), and $\varepsilon$ is zero-mean i.i.d. label noise. For the student model $f$, we follow Mei et al., (2018, Section 2.1) and parameterize a two-layer neural network with fixed second layer as:

which we train to minimize the objective (3) using PDA. Note that $\alpha=1$ corresponds to the mean field regime (which we are interested in), whereas setting $\alpha=1/2$ leads to the kernel (NTK) regimeWe use the term kernel regime only to indicate the parameter scaling $\alpha$; this does not necessarily imply that the NTK linearization is an accurate description of the trained model..

For Figure 1(a)(b) we design synthetic experiments for both regression and classification tasks, where the student model is a two-layer tanh network with $M=500$. For regression, we take the target function $f_{*}$ to be a multiple-index model with $m$ neurons: $f_{*}(x)=\frac{1}{\sqrt{m}}\sum_{i=1}^{m}\sigma_{*}(\langle w^{*}_{i},x\rangle)$, and the input is drawn from a unit Gaussian $\mathcal{N}(0,I_{p})$. For binary classification, we consider a simple two-dimensional dataset from sklearn.datasets.make$\_$circles (Pedregosa et al.,, 2011), in which the goal is to separate two groups of data on concentric circles (red and blue in Figure 1(b)). We include additional experimental results in Appendix F.

We optimize the squared loss for regression and the logistic loss for binary classification. The model is trained by PDA with batch size 50. We scale the number of inner loop steps $T_{t}$ with $t$, and the step size $\eta_{t}$ with $1/\sqrt{t}$, where $t$ is the outer loop iteration; this heuristic is consistent with the required inner-loop accuracy in Theorem 1 and Proposition 2.

2 Empirical Findings

In Figure 1(a) we verify the $O(T^{-1})$ iteration complexity of the outer loop in Theorem 1. We apply PDA to optimize the expected risk (analogous to one-pass SGD) in the regression setting, in which the input dimensionality $p=1$ and the target function is a single-index model ($m=1$) with tanh activation. We employ the resampled update (i.e., without warm-start; see Appendix B) with hyperparameters $\lambda_{1}=10^{-2},\lambda_{2}=10^{-3}$. To compute the entropy in the objective (3), we adopt the $k$-nearest neighbors estimator (Kozachenko and Leonenko,, 1987) with $k=10$.

In Figure 1(b) we visualize the evolution of neural network parameters optimized by PDA in a 2-dimensional classification problem. Due to structure of the input data (concentric rings), we expect that for a two-layer neural network to be a good separator, its parameters should also distribute on a circle. Indeed the converged solution of PDA (bright yellow) agrees with this intuition and demonstrates that PDA learns useful features beyond the kernel regime.

In Figure 1(c) we report the training and test performance of PDA in separating odd vs. even digits from the MNIST dataset. We subsample $n=2500$ training examples with binary labels, and learn a two-layer tanh network with width $M=2500$. We use the resampled update of PDA to optimize the cross entropy loss, with hyperparameters $\lambda_{1}=10^{-2},\lambda_{2}=10^{-4}$. Observe that the algorithm achieves good generalization performance (green) and roughly maintainsNote that the estimated training objective (red) slightly deviates from the ideal $1/T$-rate; this may be due to inaccuracy in the entropy estimation, or non-convergence of the algorithm (i.e., overestimation of $\mathcal{L}(q_{*})$). the $O(T^{-1})$ iteration complexity (red) in optimizing the training objective (3).

CONCLUSION

We proposed the particle dual averaging (PDA) algorithm for optimizing two-layer neural networks in the mean field regime. Leveraging tools from finite-dimensional convex optimization developed in the original dual averaging method, we established quantitative convergence rate of PDA for regularized empirical and expected risk minimization. We also provided particle complexity analysis and generalization bounds for both regression and classification problems. Our theoretical findings are aligned with experimental results on neural network optimization. Looking forward, we plan to investigate specific problem instances in which convergence rate can be obtained under vanishing regularization. It is also important to consider accelerated variants of PDA to further improve the convergence rate in the empirical risk minimization setting. Another interesting direction would be to explore other applications of PDA beyond two-layer neural networks, such as deep models (Araújo et al.,, 2019; Nguyen and Pham,, 2020; Lu et al.,, 2020; Pham and Nguyen,, 2021), as well as other optimization problems for entropic regularized nonlinear functional.

Acknowledgment

The authors would like to thank Murat A. Erdogdu and anonymous NeurIPS reviewers for their helpful feedback. AN was partially supported by JSPS Kakenhi (19K20337) and JST-PRESTO (JPMJPR1928). DW was partially supported by NSERC and LG Electronics. TS was partially supported by JSPS KAKENHI (18H03201), Japan Digital Design and JST CREST.

References

MISSING PROOFS

Then, we can show $q\propto\exp\left(-\frac{H(\theta)+\lambda_{1}\|\theta\|_{2}^{2}}{\lambda_{2}}\right)$ is an optimal solution of the problem (9) as follow. Clearly, $q\in\mathcal{P}_{2}$ and the assumption on $q$ in Lemma A with $q^{\prime}\in\mathcal{P}_{2}$ holds. Hence, for $\forall q^{\prime}\in\mathcal{P}_{2}$,

A. 2 Log-Sobolev and Talagrand’s Inequalities

The log-Sobolev inequality is useful in establishing the convergence rate of Langevin algorithm.

It is well-known that strong log-concave densities satisfy the LSI with a dimension-free constant (up to the spectral norm of the covariance).

In addition, the LSI is preserved under bounded perturbation, as originally shown in Holley and Stroock, (1987). We also provide a proof for completeness.

where we used the non-negativity of the integrand for the second inequality. ∎

We next introduce Talagrand’s inequality.

The next theorem gives a relationship between KL divergence and $2$-Wasserstein distance.

B Proof of Main Results

In this section, we prove the main theorem that provides the convergence rate of the dual averaging method. We first introduce a slight extension of PDA (Algorithm 1) which incorporates two different initializations at each outer loop step. We refer to the two versions as the warm-start and the resampled update, respectively. Note that Algorithm 1 in the main text only includes the warm-start update. In Appendix C we provide particle complexity analysis for both updates. We remark that the benefit of resampling strategy is the simplicity of estimation of approximation error $|h_{x}^{(t)}-h_{q^{(t)}}(x_{t})|$, because $h_{x}^{(t)}$ is composed of i.i.d particles and a simple concentration inequality can be applied to estimate this error.

On the other hands, the convergence analysis of warm-start scheme requires the convergence of mean field limit due to certain technical difficulties, that is, we show the convergence of Algorithm 3 with warm-start by coupling the update with its mean field limit (Algorithm 2) and taking into account the discretization error which stems from finite-particle approximation.

We now present generalized version of the outer loop convergence rate of DA. We highlight the tolerance factor $\epsilon$ in the generalized assumption (A3’) in blue.

(A2’) $|h_{\theta}(x)|\leq 1$ and $h(\theta,x)$ is smooth w.r.t. $\theta$ for $x\in\mathcal{X}$.

Under Assumptions (A1’), (A2’), and (A3’) with $\epsilon\geq 0$, for arbitrary $q_{*}\in\mathcal{P}_{2}$, iterates of the general DA method (Algorithm 4) satisfies

In the proofs, we use the following notations which are consistent with the description of Algorithm 3 and 4:

B. 2 Auxiliary Lemmas

We introduce several auxiliary results used in the proof of Theorem 1 (Theorem B) and Corollary 1. The following lemma provides a tail bound for Chi-squared variables (Laurent and Massart,, 2000).

Based on Lemma C, we get the following bound.

We set $p(\theta)=\exp(-\|\theta\|_{2}^{2}/2\sigma^{2})/Z$. Then,

where the last equality comes from the variance of Gaussian distribution.

From the symmetry of (11), and applying (11) again with (12),

From Lemma B and Example B, we see $q_{*}$ satisfies the log-Sobolev inequality with a constant $2\lambda/\exp(4c)$. As a result, $q_{*}$ satisfies Talagrand’s inequality with the same constant from Theorem A. Hence, by combining the above two inequalities, we have

where we used Pinsker’s theorem for the last inequality. This finishes the proof. ∎

where we used (12) and Gaussian integral for the last inequality. ∎

We next bound the first term in the last equation as follows.

where for the first inequality we used a similar inequality as in (12) and for the second inequality we used the Gaussian integral. Hence, we get

Recall the definition of $g^{(t)},\overline{g}^{(t)}$ and $q^{(t)}_{*}$ (see notations in subsection B. 1). We set $\gamma_{t+1}=\frac{\sum_{s=1}^{t}s}{\lambda_{2}\sum_{s=1}^{t+1}s}=\frac{t}{\lambda_{2}(t+2)}$. Note that for $t\geq 1$,

Therefore, we have for $t\geq 2$ from Proposition A with $\delta=1/t<1$,

B. 3 Outer Loop Complexity

Based on the auxiliary results and the convex optimization theory developed in Nesterov, (2009); Xiao, (2009), we now prove Theorem B which is an extension of Theorem 1.

From the definition, the density $q^{(t+1)}_{*}\in\mathcal{P}_{2}$ calculated in Algorithm 4 maximizes $V_{t}(q)$. We denote $V_{t}^{*}=V(q^{(t+1)}_{*})$. Then, for $t\geq 2$, we get

where for the first inequality we used the optimality of $q^{(t)}_{*}$ and the strong convexity (10) at $q^{(t)}_{*}$, and for the final inequality we used Lemma E.

We set $R_{t}=\left(\frac{3}{2}p+15\right)\frac{\lambda_{2}}{\lambda_{1}}\log(1+t)$ and also $\gamma_{t+1}=\frac{\sum_{s=1}^{t}s}{\lambda_{2}\sum_{s=1}^{t+1}s}=\frac{t}{\lambda_{2}(t+2)}$, as done in the proof of Lemma E.

Using (17) and applying Lemma D with $\sigma^{2}=\frac{1}{2\gamma_{t}\lambda_{1}},~{}\frac{1}{2\gamma_{t+1}\lambda_{1}}$ and $R=R_{t}$, we have for $t\geq 2$,

where for the fifth inequality we used (15) and for the sixth inequality we used $15\lambda_{2}/\lambda_{1}\leq R_{t}$.

Combining (16) and (18), we have for $t\geq 2$,

where we set $\alpha_{t}=O\left((1+\exp(8/\lambda_{2}))p^{2}\lambda_{2}\log^{2}(1+t)\right)$.

Summing the inequality (19) over $t\in\{2,\ldots,T+1\}$,

where we used $\lambda_{2}t\left|e(q^{(t)})-e(q^{(t)}_{*})\right|=\alpha_{t}$ (Lemma E), $2\alpha_{t}=O(\alpha_{t})$, and $e(q^{(T+2)}_{*})\geq 0$.

On the other hand, for $\forall q_{*}\in\mathcal{P}_{2}$,

Using (A1’), (A2’), and (A3’), we have for any density function $q$,

Hence, from (20), (21), (22), and the convexity of the loss,

Taking the expectation with respect to the history of examples, we have

B. 4 Inner Loop Complexity

We next prove Corollary 1 which gives an estimate of inner loop iteration complexity. This result is derived by utilizing the convergence rate of the Langevin algorithm under LSI developed in Vempala and Wibisono, (2019). We here consider the ideal Algorithm 2 (i.e., warm-start and exact mean field limit ($\epsilon=0$)).

We verify the assumptions required in Theorem 2. We recall that $q^{(t+1)}_{*}$ takes the form of Boltzmann distribution: for $t\geq 1$,

ADDITIONAL RESULTS AND DISCUSSIONS

For the resampling scheme, we can easily derive that the required number of particles is $O(\epsilon^{-2}\log(T/\delta))$ with high probability $1-\delta$, because i.i.d. particles are obtained by the Langevin algorithm and Hoeffding’s inequality is applicable.

Let $Z,Z_{1},\ldots,Z_{m}$ be i.i.d. random variables taking values in $[-a,a]$ for $a>0$. Then, for any $\rho>0$, we get

C. 2 Case of Warm-start

We next consider the warm-start scheme. Note that the convergence of PDA with warm-start is guaranteed by coupling it with its mean-field limit $M\rightarrow\infty$ and applying Theorem 1 without tolerance (i.e., $\epsilon=0$). To analyze the particle complexity, we make an additional assumption regarding the regularity of the loss function and the model.

(A5) $h(\cdot,x)$ is $1$-Lipschitz continuousWLOG the Lipschitz constant is set to 1, since the same analysis works for any fixed constant. for $\forall x\in\mathcal{X}$.

The above regularity assumption is common in the literature and cover many important problem settings in the optimization of two-layer neural network in the mean field regime. Indeed, (A5) is satisfied for two-layer network in Example 1 when the output or input layer is fixed and when the activation function is Lipschitz continuous.

The following proposition shows the convergence of Algorithm 1 to Algorithm 2 as $M\rightarrow\infty$.

($i$) Here we show the first statement of the proposition. We set $\rho_{1}=0$ and $\overline{\rho}_{1}=0$. We define $\rho_{t}$ and $\overline{\rho}_{t}$ recursively as follows.

Consider the inner loop at $t$-the outer step. Then, for an event where (23) holds,

we see $\overline{\rho}_{T+1}\rightarrow 0$ as $M\rightarrow+\infty$. Then, the proof is finished because for $\forall t\in\{1,\ldots,T+1\}$ and $\forall s\in\{1,\ldots,T\}$ with high probability $1-\delta$,

($ii$) We next show the second statement of the proposition. We change the definition (C. 2) of $\rho_{t+1}$ as follows:

where we used $0<1+\frac{(1-2\lambda_{1})t\eta_{t}}{\lambda_{2}(t+2)}<1$ and $\lambda_{1}\geq\frac{3}{2}$.

Noting that $(1-x)^{1/x}\leq\exp(-1)$ for $\forall x\in(0,1]$, we see that

where we used $T_{t}\geq\frac{3\lambda_{2}\log\left(4\right)}{(2\lambda_{1}-1)\eta_{t}}$. Hence, we know that for $t$,

This inequality obviously holds for $t=1$ because $\overline{\rho}_{1}=0$. We suppose it is true for $t\leq T$. Then,

Hence, the inequality (28) holds for $\forall t\in\{1,\ldots,T+1\}$, yielding

In summary, it follows that for $\forall t\in\{1,\ldots,T+1\}$ and $\forall s\in\{1,\ldots,T\}$ with high probability $1-\delta$,

where we used (26). This completes the proof. ∎

D Generalization Bounds for Empirical Risk Minimization

In this section, we give generalization bounds for the problem (3) in the context of empirical risk minimization, by using techniques developed by Chen et al., (2020). We consider the smoothed hinge loss and squared loss for binary classification and regression problems, respectively.

where $\sigma=(\sigma_{i})_{i=1}^{n}$ are i.i.d random variables taking $-1$ or $1$ with equal probability.

We introduce the uniform bound using the empirical Rademacher complexity (see Mohri et al., (2012)).

The contraction lemma (see Shalev-Shwartz and Ben-David, (2014)) is useful in estimating the Rademacher complexity.

The Rademacher complexity of this function class is obtained by Chen et al., (2020).

Suppose $|h_{\theta}(x)|\leq 1$ holds for $\forall\theta\in\Omega$ and $\forall x\in\mathcal{X}$. We have for any constant $R\leq\frac{1}{2}$ and set $S\subset\mathcal{X}$ of size $n$,

D. 2 Generalization Bound on the Binary Classification Problems

We here give a generalization bound for the binary classification problems. Hence, we suppose $\mathcal{Y}=\{-1,1\}$ and consider the problem (3) with the smoothed hinge loss defined below.

We next define the set of composite functions of loss and mean field neural networks as follows:

where we used Lemma I for the last inequality.

This theorem results in the following corollary:

Suppose the same assumptions in Theorem C hold. Moreover, we set $\lambda_{1}=\lambda/\sqrt{n}$ $(\lambda>0)$ and $\lambda_{2}=1/\sqrt{n}$. Then, the following bound holds with the probability at least $1-\delta$ over the choice of training examples,

where $p_{0}^{\prime}$ is the Gaussian distribution in proportion to $\exp(-\lambda\|\cdot\|_{2}^{2})$.

D. 3 Generalization Bound on the Regression Problem

This theorem results in the following corollary:

Suppose the same assumptions in Theorem D hold. Moreover, we set $\lambda_{1}=\lambda/\sqrt{n}$ $(\lambda>0)$ and $\lambda_{2}=1/\sqrt{n}$. Then, the following bound holds with the probability at least $1-\delta$ over the choice of training examples,

where $p_{0}^{\prime}$ is the Gaussian distribution in proportion to $\exp(-\lambda\|\cdot\|_{2}^{2})$.

E Additional Discussions

The average $\nabla_{\theta_{r}}\overline{g}^{(t)}(\Theta^{(k)})$ can then be computed as

where we use $\{\theta_{r}^{(k)}\}_{k=1}^{M}$ to denote parameters $\Theta^{(k)}$ at step $k$ of the inner loop. This formulation makes Algorithm 1 straightforward to implement.

In addition, the PDA algorithm can also be implemented with mini-batch update, in which a set of data indices $I_{t}=\{i_{t,1},\ldots,i_{t,b}\}\subset\{1,2,\ldots,n\}$ is selected per outer loop step instead of one single index $i_{t}$. Due to the reduced variance, mini-batch update can stabilize the algorithm and lead to faster convergence. Our theoretical results in the sequel trivially extends to the mini-batch setting.

E. 2 Extension to Multi-class Classification

where we supposed the outputs of $h_{\theta}$ and $h_{q}$ are also indexed by $\mathcal{C}$. Hence, the counterpart of $g^{(t)}$ in Algorithm 2 in this setting is

Using this function, the DA method for multi-class classification problems can be obtained in the same manner as Algorithm 2. Moreover, its discretization can be also immediately derived by replacing the function $\overline{g}^{(t)}$ used in Algorithm 1 with

In the case of empirical risk minimization, we can adopt an efficient implementation as done in Section E. 1. We use $\{w_{i,y}\}_{i\in\{1,\ldots,n\},y\in\mathcal{C}}$ (initialized as zeros) to store the coefficients of $h_{y}(\cdot,x_{i})$. At step $t$ in the outer loop, $w_{i_{t},y}$ ($y\in\mathcal{C}$) are updated as

Then, $\nabla_{\theta_{r}}\overline{g}^{(t)}(\Theta^{(k)})$ can be computed as

where we use $\{\theta_{r}^{(k)}\}_{k=1}^{M}$ to denote parameters $\Theta^{(k)}$ at step $k$ of the inner loop.

Finally, we remark that while we here utilize a simple network $h_{\theta}(x)$ to recover a normal two-layer neural network, it is also possible to use deep narrow networks or narrow convolutional neural networks as a component $h_{\theta}(x)$; in other words $h_{\Theta}$ can represent an ensemble of various types of small network. While such extensions are not covered by our current theoretical analysis, they may achieve better practical performance.

E. 3 Correspondence with Finite-dimensional Dual Averaging Method

We explain the correspondence between the finite-dimensional dual averaging method developed by Nesterov, (2005, 2009); Xiao, (2009) and our proposed method (Algorithm 2); our goal here is to provide an intuitive understanding of the derivation of Algorithm 2 in the context of the classical dual averaging method.

Let $\{w^{(s)}\}_{s=1}^{t}$ and $\{f^{(s)}\}_{s=1}^{t}=\{\partial_{w}l(w^{(s)},z_{s})\}_{s=1}^{t}$ be histories of iterates and stochastic gradients. The subproblems to produce the next iterate in the dual averaging method is designed by using the strongly convex function $d(w)$ and positive hyperparameters $\{\alpha_{s}\}_{s=1}^{\infty}$ and $\{\beta_{s}\}_{s=2}^{\infty}$. Specifically, the next iterate $w^{(t+1)}$ is defined as the minimizer of the following problem in which the loss function is linearized and weighted sum of which is taken over the history:

Next, we consider our problem setting of optimizing the probability distribution and reformulate the subproblem (7) solved in Algorithm 2 as follows:

F Additional Experiments

Model with the NTK scaling (green) generalizes worse than the mean field models (red and blue). This is consistent with observations in Chizat and Bach, 2018a .

For the mean field scaling, PDA (under early stopping) leads to slightly lower test error than noisy GD. We intend to further investigate this difference in the generalization performance. (see Appendix D for generalization bounds of the PDA solution)

Figure 3(a) demonstrates the advantage of $L_{p}$-norm regularization for $p<2$ in empirical risk minimization, when the target function exhibits sparse structure. We set $n=1000,p=50$; the teacher is a multiple-index model ($m=2$) with binary activation, and parameters of each neuron are $1$-sparse. We optimize the student model with PDA (warm-start), where we set $\lambda_{1}=10^{-2}$, $\lambda_{2}=10^{-4}$, and vary the norm penalty $p$ from 1.01 to 2. Note that smaller $p$ results in favorable generalization due to the induced sparsity. On the other hand, we expect the benefit of sparse regularization to diminish when the target function is not sparse. This intuition is confirmed in 3(b), where we control the target sparsity by randomly selecting $r$ parameters to be non-zero, and we define $s=r/d$ to be the sparsity level. Observe that the benefit of sparsity-inducing regularization (smaller $p$) is more prominent under small $s$ (brighter color), which indicates a sparse target function.

F. 3 On the Role of Entropy Regularization

F. 4 Adaptivity of Mean Field Neural Networks

Recall that one motivation to study the mean field regime (instead of the kernel regime) is the presence of feature learning. We illustrate this behavior in a simple student-teacher setup, where the target function is a single-index model with tanh activation. We set $n=500,d=50$, and optimize a two-layer tanh network ($M=1000$), either in the mean field regime using PDA, or in the kernel regime using SGD. For both methods we choose $\lambda_{1}=10^{-3}$, and for PDA we choose $\lambda_{2}=10^{-4}$.

In Figure 5 we plot the the evolution of the cosine similarity between the target vector $w^{*}$ and the top-5 singular vectors (PC1-5) of the weight matrix during training. In Figure 5(a) we observe that the mean field model trained with PDA “adapts” to the low-dimensional structure of the target function; in particular, the leading singular vector (bright yellow) aligns with the target direction. In contrast, we do not observe such alignment on the network in the kernel regime (Figure 5(b)), because the parameters do not travel away from the initialization. This comparison demonstrates the benefit of the mean field parameterization.

G Additional Related Work

Bayesian inference is another example distribution optimization, in which the objective is to minimize an entropic regularized linear functional. In addition to the Langevin algorithm, several interacting particle methods have been developed for this purpose, such as particle mirror descent (PMD) (Dai et al.,, 2016), Stein variational gradient descent (SVGD) (Liu and Wang,, 2016), and ensemble Kalman sampler (Garbuno-Inigo et al.,, 2020), and the corresponding mean field limits have been analyzed in Lu et al., (2019); Ding and Li, (2019). We remark that naive gradient-based method on the probability space often involves computing the probability of particles for the entropy term (e.g., kernel density estimation in PMD), which presents significant difficulty in constructing particle inference algorithms. In contrast, our proposed algorithm avoids this computational challenge due to its algorithmic structure.

Parallel to our work, several recent papers also proposed optimization methods over space of probability measures by adapting finite-dimensional convex optimization theory. Ying, (2020), Kent et al., (2021) and Chizat, (2021) extend the Mirror descent method, Frank-Wolfe method, and (accelerated) Bregman proximal gradient method to the optimization of probability measures, respectively. In addition, Hsieh et al., (2019) developed an entropic mirror descent algorithm for generative adversarial networks, and Chu et al., (2019) analyzed probability functional descent in the context of variational inference and reinforcement learning.

The neural tangent kernel model (Jacot et al.,, 2018) describes the learning dynamics of neural network under appropriate scaling. Such description builds upon the linearization of the learning dynamics around its initialization, and (quantitative) global convergence guarantees of gradient-based methods for neural networks can be shown for regression problems (Du et al.,, 2019; Allen-Zhu et al.,, 2019; Zou et al.,, 2020; Nitanda and Suzuki,, 2021) as well as classification problems (Cao and Gu,, 2019; Nitanda et al.,, 2019; Ji and Telgarsky,, 2019).

However, due to the linearization, the NTK model cannot explain the presence of “feature learning” in neural networks (i.e. parameters are able to travel and adapt to the structure of the learning problem). In fact, various works have shown that deep learning is more powerful than kernel methods in terms of approximation and estimation error (Suzuki,, 2019; Ghorbani et al., 2019b, ; Suzuki and Nitanda,, 2021; Schmidt-Hieber,, 2020; Ghorbani et al.,, 2020; Imaizumi and Fukumizu,, 2020), and in certain settings, neural networks optimized with gradient-based methods can outperform the NTK model (or more generally any kernel methods) in terms of generalization error or excess risk (Allen-Zhu and Li,, 2019; Ghorbani et al., 2019a, ; Yehudai and Shamir,, 2019; Bai and Lee,, 2019; Allen-Zhu and Li,, 2020; Li et al.,, 2020; Suzuki,, 2020; Daniely and Malach,, 2020).