Particle Dual Averaging: Optimization of Mean Field Neural Network with Global Convergence Rate Analysis
Year: 2022

Abstract:
We propose the particle dual averaging (PDA) method, which generalizes the dual averaging method in convex optimization to the optimization over probability distributions with quantitative runtime guarantee. The algorithm consists of an inner loop and outer loop: the inner loop utilizes the Langevin algorithm to approximately solve for a stationary distribution, which is then optimized in the outer loop. The method can thus be interpreted as an extension of the Langevin algorithm to naturally handle nonlinear functional on the probability space. An important application of the proposed method is the optimization of neural network in the mean field regime, which is theoretically attractive due to the presence of nonlinear feature learning, but quantitative convergence rate can be challenging to obtain. By adapting finite-dimensional convex optimization theory into the space of measures, we analyze PDA in regularized empirical / expected risk minimization, and establish quantitative global convergence in learning two-layer mean field neural networks under more general settings. Our theoretical results are supported by numerical simulations on neural networks with reasonable size.

heading: 1 Introduction
text: 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, 2018; 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),
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
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.

heading: 1.1 Contributions
text: 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). The resulting algorithm has comparable per-iteration cost as gradient descent and can be efficiently implemented.
For optimizing two-layer neural network in the mean-field regime, we establish quantitative global convergence rate of PDA in minimizing an KL-regularized objective: the algorithm requires Õ(✏ 3) steps and Õ(✏ 2) particles to reach an ✏-accurate solution, where Õ hides logarithmic factors. Importantly, our analysis does not couple the learning dynamics with certain continuous time limit, but directly handles the discrete update. This leads to a simpler analysis that covers more general settings. We also derive the generalization bound on the solution obtained by the algorithm. From the viewpoint of the optimization, PDA is an extension of Langevin algorithm to handle entropicregularized nonlinear functionals on the probability space. Hence we believe our proposed method can also be applied to other distribution optimization problems beyond the training of neural networks.

heading: 1.2 Related Literature
text: Mean field limit of two-layer NNs. 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., 2006; 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.
Javanmard et al. (2019) analyzed a particular RBF network and established linear convergence (up to certain error1) for strongly concave target functions. Rotskoff et al. (2019) provided a sublinear rate in continuous time for a modified gradient flow. In the regularized setting, Chizat (2019) obtained local linear convergence under certain non-degeneracy assumption on the objective. Wei et al. (2019) also proved polynomial rate for a perturbed dynamics under weak `2 regularization.
1Note that such error yields sublinear rate with respect to arbitrarily small accuracy ✏.
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 algorithm. 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.

heading: 1.3 Notations
text: Let R+ denote the set of non-negative real numbers and k · k2 the Euclidean norm. Given a density function q : Rp ! R+, we denote the expectation with respect to q(✓)d✓ by Eq[·]. For a function f : Rp ! R, we define Eq[f ] = R f(✓)q(✓)d✓ when f is integrable. KL is the Kullback-Leibler divergence: KL(qkq0) def= R q(✓) log ⇣ q(✓) q0(✓) ⌘ d✓. Let P2 denote the set of positive densities q on Rp such that the second-order moment Eq[k✓k22] <1 and entropy 1 < Eq[log(q)] < +1 are well defined. N (0, Ip) is the Gaussian distribution on Rp with mean 0 and covariance matrix Ip.

heading: 2 Problem Setting
text: We consider the problem of risk minimization with neural networks in the mean field regime. For simplicity, we focus on supervised learning. We here formalize the problem setting and models. Let X ⇢ Rd and Y ⇢ R be the input and output spaces, respectively. For given input data x 2 X , we predict a corresponding output y = h(x) 2 Y through a hypothesis function h : X ! Y .

heading: 2.1 Neural Network and Mean Field Limit
text: We adopt a neural network in the mean field regime as a hypothesis function. Let ⌦ = Rp be a parameter space and h✓ : X ! Y (✓ 2 ⌦) be a bounded function which will be a component of a neural network. We sometimes denote h(✓, x) = h✓(x). Let q(✓)d✓ be a probability distribution on the parameter space ⌦ and ⇥ = {✓r}Mr=1 be the set of parameters ✓r sampled from q(✓)d✓. A hypothesis is defined as an ensemble of h✓r as follows:
h⇥(x) def =
1
M
MX
r=1
h✓r (x). (1)
A typical example in the literature of the above formulation is a two-layer neural network. Example 1 (Two-layer Network). Let ar 2 R and br 2 Rd (r 2 {1, 2, . . . ,M}) be parameters for output and input layers, respectively. We set ✓r = (ar, br) and ⇥ = {✓r}Mr=1. Denote h✓r (x) def =
2(ar 1(b>r x)) (x 2 X ), where 1 and 2 are smooth activation functions. Then the hypothesis h⇥ is a two-layer neural network composed of neurons h✓r : h⇥(x) = 1 M P M r=1 2(ar 1(b > r x)).
Remark. The purpose of 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⇥ of deep neural networks h✓r .
Suppose the parameters ✓r follow a probability distribution q(✓)d✓, then h⇥ can be viewed as a finite-particle discretization of the following expectation,
hq(x) = Eq[h✓(x)]. (2)
which we refer to as the mean field limit of the neural network h⇥. As previously discussed, when h⇥ is overparameterized, optimizing h⇥ becomes “close” to directly optimizing the probability distribution on the parameter space ⌦, for which convergence to the optimal solution may be established under appropriate conditions (Nitanda and Suzuki, 2017; Mei et al., 2018; Chizat and Bach, 2018b). Hence, the study of optimization of hq with respect to the probability distribution q(✓)d✓ may shed light on important properties of overparameterized neural networks.

heading: 2.2 Regularized Empirical Risk Minimization
text: We briefly outline our setting for regularized expected / empirical risk minimization. The prediction error of a hypothesis is measured by the loss function `(z, y) (z, y 2 Y), such as the squared loss `(z, y) = 0.5(z y)2 for regression, or the logistic loss `(z, y) = log(1 + exp( yz)) for binary classification. Let D be a data distribution over X ⇥ Y . For expected risk minimization, the distribution D is set to the true data distribution; whereas for empirical risk minimization, we take D to be the empirical distribution defined by training data {(xi, yi)}ni=1 (xi 2 X , yi 2 Y) independently sampled from the data distribution. We aim to minimize the expected / empirical risk together with a regularization term, which controls the model complexity and also stabilizes the optimization. The regularized objective can be written as follows: for 1, 2 > 0,
min q2P2
n L(q) def= E(X,Y )⇠D[`(hq(X), Y )] +R 1, 2(q) o , (3)
where R 1, 2 is a regularization term composed of the weighted sum of the second-order moment and negative entropy with regularization parameters 1, 2:
R 1, 2(q) def = 1Eq[k✓k22] + 2Eq[log(q)]. (4)
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.

heading: 2.3 The Langevin Algorithm
text: Before presenting our proposed method, we briefly review the Langevin algorithm. For a given smooth potential function f : ⌦! R, the Langevin algorithm performs the following update: given the initial ✓(1) ⇠ q(1)(✓)d✓, step size ⌘ > 0, and Gaussian noise ⇣(k) ⇠ N (0, Ip),
✓(k+1) ✓(k) ⌘r✓f(✓(k)) + p 2⌘⇣(k). (5)
Under appropriate conditions on f , it is known that ✓(t) converges to a stationary distribution proportional to exp( f(·)) in terms of KL divergence at a linear rate (e.g., Vempala and Wibisono (2019)) up to O(⌘)-error, where we hide additional factors in the big-O notation. Alternatively, note that when the normalization constant R exp( f(✓))d✓ exists, the Boltzmann distribution in proportion to exp( f(·)) is the solution of the following optimization problem,
min q:density {Eq[f ] + Eq[log(q)]} . (6)
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.
However, the objective (3) that we aim to optimize is beyond the scope of Langevin algorithm – due to the nonlinearity of loss `(z, y) with respect to z, the stationary distribution cannot be described as a closed-form solution of (6). To overcome this limitation, we develop the particle dual averaging (PDA) algorithm which efficiently solves (3) with quantitative runtime guarantees.

heading: 3 Proposed Method
text: 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.

heading: 3.1 Particle Dual Averaging
text: 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.
• In the outer loop, the last iterate ⇥̃(t) of the previous inner loop is given. We compute @z`(h⇥̃(t)(xt), yt), which is a component of the Gibbs potential
2, and initialize a set of particles ⇥(1) at ⇥̃(t). In Appendix B we introduce a different “restarting” scheme for the initialization.
• In the inner loop, we run the Langevin algorithm (noisy gradient descent) starting from ⇥(1), where the gradient at the k-th inner step is given byr✓g(t)(✓(k)r ), which is a sum of weighted average of @z`(h⇥̃(s)(xs), ys)@✓h(✓ (k) r , xs) and the gradient of `2-regularization (see Algorithm 1).
Algorithm 1 Particle Dual Averaging (PDA)
Input: data distribution D, initial density q(1), number of outer-iterations T , learning rates {⌘t}Tt=1, number of inner-iterations {Tt}Tt=1 Randomly draw i.i.d. initial parameters ✓̃(1)r ⇠ q(1)(✓)d✓ (r 2 {1, 2, . . . ,M}) ⇥̃(1) {✓̃(1)r }Mr=1 for t = 1 to T do
Randomly draw data (xt, yt) from D ⇥(1) = {✓(1)r }Mr=1 ⇥̃(t) for k = 1 to Tt do Run inexact noisy gradient descent for r 2 {1, 2, . . . ,M} r✓g(t)(✓(k)r ) 2 2(t+2)(t+1) P t s=1 s@z`(h⇥̃(s)(xs), ys)@✓h(✓ (k) r , xs) + 2 1t 2(t+2) ✓(k)r
✓(k+1)r ✓(k)r ⌘tr✓g(t)(✓(k)r ) + p 2⌘t⇣ (k) r (i.i.d. Gaussian noise ⇣ (k) r ⇠ N (0, Ip))
end for ⇥̃(t+1) ⇥(Tt+1) = {✓(Tt+1)r }Mr=1
end for Randomly pick up t 2 {2, 3, . . . , T +1} following the probability P[t] = 2t
T (T+3) and return h⇥̃(t)
Figure 1 provides a pictorial illustration of Algorithm 1. Note that this procedure is a slight modification of the normal gradient descent algorithm: the first term of r✓g(t) is similar to the gradient of the loss @✓r`(h⇥(k)(x), y) ⇠ @z`(h⇥(k)(x), y)@✓h(✓ (k) r , x) where ⇥(k) = {✓(k)r }Mr=1. Indeed, if we
2In Algorithm 1, the terms @z`(h⇥̃(s)(xs), ys) appear in inner loop; but note that these terms only need to be computed in outer loop because they are independent to the inner loop iterates.
set the number of inner-iterations Tt = 1 and replace the direction r✓g(t)(✓(k)r ) with the gradient of the L2-regularized loss, then PDA exactly reduces to the standard noisy gradient descent algorithm considered in Mei et al. (2018). Algorithm 1 can be extended to the minibatch variant in the obvious manner; for efficient implementation in the empirical risk minimization setting see Appendix E. 1.

heading: 3.2 Mean Field View of PDA
text: 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:
g(t)(✓) = 2
2(t+ 2)(t+ 1)
tX
s=1
s @z`(h⇥̃(s)(xs), ys)h(✓, xs) + 1k✓k 2 2 .
Due to the rapid convergence of Langevin algorithm outlined in Subsection 2.3, the particles ✓(k+1)r (r 2 {1, . . . ,M}) can be regarded as (approximate) samples from the Boltzmann distribution: exp g(t) . 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 ! 1 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 Dual Averaging (DA)
Input: data distribution D and initial density q(1)
for t = 1 to T do Randomly draw a data (xt, yt) from D g(t) @z`(hq(t)(xt), yt)h(·, xt) + 1k · k22 Obtain an approximation q(t+1) of the density function q(t+1)⇤ / exp ⇣ Pt s=1 2sg (s)
2(t+2)(t+1)
⌘
end for Randomly pick up t 2 {2, 3, . . . , T + 1} following the probability P[t] = 2t
T (T+3) and return hq(t)
Algorithm 2 iteratively updates the density function q(t+1)⇤ 2 P2 which is a solution to the objective:
min q2P2
( Eq h tX
s=1
sg(s) i +
2 2 (t+ 2)(t+ 1)Eq[log(q)]
) , (7)
where the function g(t) = @z`(hq(t)(xt), yt)h(·, xt) + 1k · k22 is the functional derivative of `(hq(xii), yt) + 1Eq[k✓k22] with respect to q at q(t). In other words, the objective (7) is the sum of weighted average of linear approximations of loss function and the entropic regularization in the space of probability distributions. In this sense, the DA method can be seen as an extension of the Langevin algorithm to handle entropic regularized nonlinear functionals on the probability space by iteratively linearizing 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 ! 1. 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.

heading: 4 Convergence Analysis
text: 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.

heading: 4.1 Outer Loop Complexity
text: We first analyze the convergence rate of the dual averaging (DA) method (Algorithm 2). Our analysis will be made under the following assumptions. Assumption 1.
(A1) Y ⇢ [ 1, 1]. `(z, y) is a smooth convex function w.r.t. z and |@z`(z, y)|  2 for y, z 2 Y . (A2) |h(✓, x)|  1 and h(✓, x) is smooth with respect to ✓ for x 2 X .
(A3) KL(q(t+1)kq(t+1)⇤ )  1/t2.
Remark. (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 Y ⇢ [ 1, 1] and |h✓(x)|  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 2 P2,
e(q) def = Eq[log(q)]
4 2 p 2
✓ exp ✓ 4
2
◆ + log ✓ 3⇡ 2 1 ◆◆ .
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)⇤ ) 0. We have the following convergence rate of DA3. Theorem 1 (Convergence of DA). Under Assumptions (A1), (A2), and (A3), for arbitrary q⇤ 2 P2, iterates of the DA method (Algorithm 2) satisfies
2
T (T + 3)
T+1X
t=2
t ⇣ E[L(q(t))] L(q⇤) ⌘
 O ⇣ 1 T 2 1 + 1Eq⇤ ⇥ k✓k22 ⇤ + 2e(q⇤) T + 2 T (1 + exp(8/ 2))p 2 log2(T + 2) ⌘ ,
where the expectation E[L(q(t))] is taken with respect to the history of examples.
Theorem 1 demonstrates the convergence rate of Algorithm 2 to the optimal value of the regularized objective (3) in expectation. Note that 2
T (T+3) P T+1 t=2 tE[L(q(t))] is the expectation of E[L(q(t))]
according to the probability P[t] = 2t T (T+3) (t 2 {2, . . . , T + 1}) as specified in Algorithm 2. If we take p, 1, 2 as constants and use Õ to hide the logarithmic terms, we can deduce that after Õ(✏ 1) iterations, an ✏-accurate solution of the optimal distribution: L(q)  infq2P2 L(q) + ✏ is achieved in expectation. Importantly, this convergence rate applies to any choice of regularization parameters, in contrast to the strong regularization required in Hu et al. (2019); Jabir et al. (2019).
On the other hand, due to the exponential dependence on 12 , our convergence rate is not informative under weak regularization 2 ! 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.
3In Appendix B we introduce a more general version of Theorem 1 that allows for inexact hq(t)(x), which simplifies the analysis of finite-particle discretization presented in Appendix C.

heading: 4.2 Inner Loop Complexity
text: 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): Theorem 2 (Vempala and Wibisono (2019)). Consider a probability density q(✓) / exp( f(✓)) satisfying the log-Sobolev inequality with constant ↵, and assume f is smooth andrf is L-Lipschitz, i.e., kr✓f(✓) r✓f(✓0)k2  Lk✓ ✓0k2. If we run the Langevin algorithm (5) with learning rate 0 < ⌘  ↵4L2 and let q (k)(✓)d✓ be a probability distribution that ✓(k) follows, then we have,
KL(q(k)kq)  exp( ↵⌘k)KL(q(1)kq) + 8↵ 1⌘pL2.
Theorem 2 implies that a -accurate solution in KL divergence can be obtained by the Langevin algorithm with ⌘  ↵4L2 min n 1, 4p o and 1 ↵⌘ log 2KL(q (1)kq) -iterations.
Since the optimal solution of a sub-problem in DA (Algorithm 2) takes the forms of q(t+1)⇤ / exp ⇣ Pt s=1 2sg (s)
2(t+2)(t+1)
⌘ , we can verify the LSI and determine the constant for q(t+1)⇤ (✓)d✓ based on
the LSI perturbation lemma from Holley and Stroock (1987) (see Lemma B and Example 2 in Appendix A. 2). Consequently, we can apply Theorem 2 to q(t+1)⇤ for the inner loop complexity when r✓ log q(t+1)⇤ is Lipschitz continuous, which motivates us to introduce the following assumption. Assumption 2.
(A4) @✓h(·, x) is 1-Lipschitz continuous: k@✓h(✓, x) @✓h(✓0, x)k2  k✓ ✓0k2, 8x 2 X , ✓, ✓0 2 ⌦.
Remark. (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 X is compact. We remark that this assumption can be relaxed to Hölder continuity of @✓h(·, x) via the recent result of Erdogdu and Hosseinzadeh (2020), which allows us to extend Theorem 1 to general Lp-norm regularizer for p > 1. For now we work with (A4) for simplicity of the presentation and proof.
Set t+1 to be the desired accuracy of an approximate solution q(t+1) specified in (A3): t+1 = 1/(t+ 1)2, we have the following guarantee for the inner loop. Corollary 1 (Inner Loop Complexity). Under (A1), (A2), and (A4), if we run the Langevin algorithm with step size ⌘t = O ⇣ 1 2 t+1
p(1+ 1)2 exp(8/ 2)
⌘ on (7), then an approximate solution satisfying
KL(q(t+1)kq(t+1)⇤ )  t+1 can be obtained within O ⇣
2 exp(8/ 2) 1⌘t log 2KL(q (t)kq(t+1)⇤ ) t+1
⌘ -iterations.
Moreover, KL(q(t)kq(t+1)⇤ ) (t 2 {1, 2, . . . , T + 1}) are uniformly bounded with respect to t as long as q(1) is a Gaussian distribution and (A3) is satisfied.
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.

heading: 4.3 Total Complexity
text: Combining Theorem 1 and Corollary 1, we can now derive the total complexity of our proposed algorithm. For simplicity, we take p, 1, 2 as constants and hide logarithmic terms in Õ and ⇥̃. The following corollary establishes a Õ(✏ 3) total iteration complexity to obtain an ✏-accurate solution in expectation because Tt = ⇥̃(t2) = Õ(✏ 2) for t  T . Corollary 2 (Total Complexity). Let ✏ > 0 be an arbitrary desired accuracy and q(1) be a Gaussian distribution. Under assumptions (A1), (A2), (A3), and (A4), if we run Algorithm 2 for T = ⇥̃(✏ 1) iterations on the outer loop, and the Langevin algorithm with step size ⌘t = ⇥ ⇣ 1 2 t+1
p(1+ 1)2 exp(8/ 2)
⌘
for Tt = ⇥̃(⌘ 1t ) iterations on the inner loop, then an ✏-accurate solution: L(q)  infq2P2 L(q) + ✏ of the objective (3) is achieved in expectation.
Quantitative convergence guarantee. To translate the above convergence rate result to the finiteparticle PDA (Algorithm 1), we also characterize the finite-particle discretization error in Appendix C. For the particle complexity analysis, we consider two versions of particle update: (i) the warmstart scheme described in Algorithm 1, in which ⇥(1) is initialized at the last iterate ⇥̃(t) of the previous inner loop, and (ii) the resampling scheme, in which ⇥(1) is initialized from the initial distribution q(1)(✓)d✓ (see Appendix B for details). Remarkably, for the resampling scheme, we provide convergence rate guarantee in time- and space-discretized settings that is polynomial in both the iterations and particle size; specifically, the particle complexity of Õ(✏ 2), together with the total iteration complexity of Õ(✏ 3), suffices to obtain an ✏-accurate solution to the objective (3) (see Appendix B and C for precise statement).

heading: 5 Experiments

heading: 5.1 Experiment Setup
text: We employ our proposed algorithm in both synthetic student-teacher settings (see Figure 2(a)(b)) and real-world dataset (see Figure 2(c)). For the student-teacher setup, the labels are generated as yi = f⇤(xi) + "i, where f⇤ is the teacher model (target function), and " 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:
f(x) = 1
M↵
MX
r=1
(w> r x+ br), (8)
which we train to minimize the objective (3) using PDA. Note that ↵ = 1 corresponds to the mean field regime (which we are interested in), whereas setting ↵ = 1/2 leads to the kernel (NTK) regime4.
Synthetic student-teacher setting. For Figure 2(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) = 1p m P m
i=1 ⇤(hw⇤i , xi), and the input is drawn from a unit Gaussian N (0, Ip). 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 2(b)). We include additional experimental results in Appendix F.
PDA hyperparameters. 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 Tt with t, and the step size ⌘t with 1/ p t, where t is the outer loop iteration; this heuristic is consistent with the required inner-loop accuracy in Theorem 1 and Proposition 2.
4We use the term kernel regime only to indicate the parameter scaling ↵; this does not necessarily imply that the NTK linearization is an accurate description of the trained model.

heading: 5.2 Empirical Findings
text: Convergence rate. In Figure 2(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 1 = 10 2, 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.
Presence of feature learning. In Figure 2(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.
Binary classification on MNIST. In Figure 2(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 1 = 10 2, 2 = 10 4. Observe that the algorithm achieves good generalization performance (green) and roughly maintains5 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 (Araú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.

heading: Acknowledgment
text: 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.

Review summaries:
In this paper the authors  propose the particle dual averaging (PDA) method, which generalizes the dual averaging method in convex optimization to the optimization over probability distributions with quantitative runtime guarantee. The method in thi paper can thus be interpreted as an extension of the Langevin algorithm to naturally handle nonlinear functional on the probability space. An important application of the proposed method is the optimization of neural network in the mean field regime. By adapting finite-dimensional convex optimization theory into the space of distributions, the athors  establish global convergence of PDA for two-layer mean field neural networks under more general settings and simpler analysis. Numerics support the   theoretical findings. 

The paper studies the convergence of the dual averaging method for the 1-hidden layer neural network in the mean-filed regime. Mean-field perspective allows one to recast original non-convex optimization problem over parameters of neural networks into (non necessarily strictly) convex problem in the space of probability measures.   Similarly to Mei et al. 2018 and Hu et al. 2019, authors consider entropy regularised cases under which the optimal solution to the infinite-dimensional problem over measures admits convenient Boltzman/Gibbs distribution.  Unlike Mei et al. 2018 and Hu et al. 2019,  who studied corresponding Langevin dynamics, while here authors study dual averaging methods in which one approximates the optimal sampling distribution by a sequence of simpler distributions that converge to it (with a rate one over a number of iterations), and use Langevin dynamics to sample from each distribution in the sequence.  If one would be able to quantify the rate of convergence of to each distribution in the sequence, this will imply the overall rate and complexity of the algorithm. 

The paper is motivated by optimizing two-layer neural networks, with entropy and quadratic regularization. The proposed method is an extension of the dual averaging method to the space of probability densities, together with a finite particle approximation based on Langevin dynamics. An error bound is provided, such that to reach an error $\epsilon$, the number of steps and the number of particles required are both polynomial in $1/\epsilon$. 

This paper introduces a Particle Dual Averaging (PDA) algorithm to minimize non-linear functional of probability distributions. One of the main application of this algorithm is the training of two-layer neural networks. The authors obtain quantitative convergence results for PDA under mild assumption on the neural network for the minimization of an entropy-regularized loss function. In particular they show that a solution can be reached with precision $O(\varepsilon)$ after $O(\varepsilon^{-3})$ iterations. The theoretical study is completed with the investigation of the finite particle algorithm and generalization bounds. The authors then present a short experimental study on regression and classification tasks, illustrating the properties of the algorithm.