Particle Dual Averaging: Optimization of Mean Field Neural Network with Global Convergence Rate Analysis
Authors: Atsushi Nitanda
Denny Wu
Taiji Suzuki
Venue: NeurIPS 2021 Poster (2021)
Keywords: Neural Network Optimization
Mean field Regime
Overparameterization
Langevin Algorithm

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 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 not only establish global convergence of PDA for two-layer mean field neural networks under more general settings and simpler analysis, but also provide quantitative polynomial runtime guarantee. Our theoretical results are supported by numerical simulations on neural networks with reasonable size. 

Particle Dual Averaging: Optimization of Mean Field Neural Network with Global Convergence Rate Analysis
Atsushi Nitanda
Denny Wu University of Toronto Vector Institute for Artificial Intelligence dennywuêca.
Kyushu Institute of Technology
RIKEN Center for Advanced Intelligence Project ni nitanda@ai kkyutccc........................ jc.jp
htwpyy/ussa..ronttoooo toronto.
Taiji Suzuki
University of Tokyo
RIKEN Center for Advanced Intelligence Project taiji@mist i.t.i.utuae u-tokyo jc.jp
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 mean field regime, which is theoretically attractive due to the presence of nonlinear feature learning, but quantitative convergence rate can challenging to obtain. By adapting finite snite-dimmensionll 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 neural networks with reasonable size.
neural network the
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 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 kernel regression with respect to the neural tangent kernel (Jacot al., 2018), which leads to global convergence at linear rate (Du et al., 2019; Allen -Zu 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, 201 (a). Such regime thus fails to explain the feature learning of neural networks (Yang and Hu, 2020), which 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 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 (NcurIPS 2021).

---

3
Initialized q
Optimal
Intermediate q!
inner loop
ID visualization of parameter distribution of mean field two-layer neu- ral network (tanh) optimized by PDA. The inner loop uses the Langevin algorithm to solve an approximate stationary distribu- tion ( which then optimized the outer loop towards the true target
Figure 1:
Density
2
Langevin iterates
inner loop
1
outer loop
which captures the potentially nonlinear evolution 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-ting time 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, important research problem that we aim 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 an important application, establish polynomial runtime guarantee for the discrete-time algorithm; to our knowledge this is the first quantitative global convergence result under similar settings.
1.1 Contributions
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 Boltzmann distribution, samples from which can obtained from the Langevin algorithm (see Figure 1). The resulting algorithm has comparable per- perriteration cost gradient descent and can be efficiently implemented.
For optimizing two-layer neural network the mean-field regime, we establish quantitative global convergence rate of PDA minimizing an KL-regularized objective: the algorithm requires Ō(c-3) steps and Ö(E (() particles to reach an e-accurate solution, where Ô 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 entropic- regularized nonlinear functionals the probability space. Hence we believe our proposed method can also be applied to other distribution optimization problems beyond the training of neural networks.
1.2 Related Literature
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 dell-deseribed by nonlinear partial differential equation (PDE), which can be viewed as solving infiniteddimension convex problem (Bengio et al., 2006; Bach, 2017). Global convergence can be derived by studying the limiting PDE (Mei al., 2018; Chizat and Bach, 2018b; Rotskoff and Vanden-Eijndenn, 2018; Sirignano and Spiliopoulos, 2020), yet quantitative convergence rate generally requires additional assumptions.
Javanmard et al. (2019) analyzed particular RBF network and established linear convergence (up to certain error') ) for strongly concave target functions. Rotskoff al. (2019) provided sublinear rate in continuous time for a modified gradient flow. In the regularized setting, Chizat (2019) obtained local linear convergence under certain non-degenyyacy assumption on the objective. Wei al. (2019) also proved polynomial rate for perturbed dynamics under weak & regularization.
Note that such error yields sublinear rate with respect arbitrarily small accuracy

---

Our setting is most related to Hu et al. (2019), who studied the minimization of nonlinear functional with KL regularization on the probability space, and showed linear convergence (in continuous time) of particle dynamics named mean field Langevin dynamics when the regularization is sufficiently strong. Chen et al. (2020) also considered optimizing KI.-regularized objective the infinite- width and continuous-time limit, and derived NTKKliike 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 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-ting 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 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. How- ever, such result often depends on the spectral gap that grows exponentially in dimensionality, which renders the analysis ineffective for neural net optimization in the eigh-dimensional parameter space.
Very recently, convergence of Hamiltonian Monte Carlo in learning certain mean field models has been analyzed in Bou- 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.
1.3 Notations
Let R++ denote the set of non-negative real numbers and || ll2 the Euclidean norm. Given a density
function q : I2 + + we denote the expectation with respect to (0) by E,I:]. For a function f: R - R, we define E., [f] = (()) ()) de when f is integrable. KL is the Kullback-Leibler
divergence: (g(gll4) def (0) log (0(
d0 Let P2 denote the set of positive densities q on RR" such that the second-order moment [1 < do and entropy -oo < - Ey (lg(()) < +00 are well defined. (0, Ip) is the Gaussian distribution on E. with mean 0 and covariance matrix Ip.
(())
2 Problem Setting
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' C Rd and y C IR be the input and output spaces, respectively. For given input data x E X, we predict a corresponding output y h(x) E y through hypothesis function h X - V.
2.1 Neural Network and Mean Field Limit
We adopt a neural network in the mean field regime as a hypothesis function. Let S = TR be a parameter space and hg X Y' (0 E 52) be bounded function which will be component of neural network. We sometimes denote h0.x)) -)= = ho(r). Let ( be a probability distribution on the parameter space Ω and θ {0, =-= be the set of parameters 0, sampled from (0))O. A hypothesis is defined as an ensemble of ha, as follows:
hz(x) def
M Σ ho, (z).
typical example in the literature of the above formulation is two-layer neural network.
Example 1 (Two-layer Network). Let a, E IR and by E Re (r E {1, 2,,, M)) be parameters for output and input layers, respectively. We set 0, = (ar, b,) and θ = [M Denote ha, (x) def az(a,o (b, x)) (x E X), where as and 02 are smooth activation functions. funictionn.. Then the hypothesis he is two-layer neural network composed of neurons ha.: he hz(z) 1/4 ΣM' o2(a,00(6)) (b,.z)). x)).
3

---

Remark. The purpose of d2 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 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 he of deep neural networks he,-
Suppose the parameters 0, follow a probability distribution g(0)dO, then he can be viewed as finite-particle discretization of the following expectation,
h1( = (xr)
which we refer to as the mean field limit of the neural network he. As previously discussed, when he overparameterized, optimizing he becomes "close" to directly optimizing the probability distribution on the parameter space Ω2 for which convergence to the optimal solution may be established under appropriate conditions (Nitanda and Suzuki, 2017; Mei et al., 2018; Chizat and Bach, 201 8b). Hence, the study of optimization of hg with respect to the probability distribution q())d8 may shed light on important properties of overparameterized neural networks.
2.2 Regularized Empirical Risk Minimization
We briefly outline our setting for regularized expected / empirical risk minimization. The prediction error of hypothesis is measured by the loss function ((z,)) v) (z, y y), such the squared loss f(z, v) = 0.5(= - y)² for regression, or the logistic loss {(z, y) = log( + exp(-yz)) for binary classification. Let D be data distribution over X' V. 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 {(x;, ):)) (ax E X. )) independently sampled from the data distribution. We aim to minimize the expected empirical risk together with regularization term, which controls the model complexity and also stabilizes the optimization. The regularized objective can be written as follows: for d1. A2 > 0,
pEP; { (Ca) del Bxx,r))D(((hh( (X),
min
Y)] (1)
where Rx,,12 is regularization term composed the weighted sum of the second-order moment and negative entropy with regularization parameters '1, A2:
RA,,Aa (g)
def
X FE
[101 + (oggg))
(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 g.¹) specified in Section 3 are well defined.
While our primary focus is the optimization of the objective (3), we can also derive generalization error bound for the empirical risk minimizer of order of O 0n211/2) for both the regression and binary classification settings, following Chen et al. (2020). We defer the details to Appendix D.
2.3 The Langevin Algorithm
Before presenting proposed method, we briefly review the Langevin algorithm. For given smooth potential function f Ω + I, the Langevin algorithm performs the following update: given the initial ((1) ~ ( c00, step size 7 > 0, and Gaussian noise ((k) (O, Ip).
(++++) +- k)*) () + (2))
(5)
Under appropriate conditions on f. is known that (t9) converges to stationary distribution proportional to exp(-f(-)) in terms of KL divergence at linear rate (e.g.. Vempala and Wibisono (2019)) up to (n))erroo,, where we hide additional factors in the big-O notation.
Alternatively, note that when the normalization constant f (xx(-f(())) exists, the Boltzmann distribution in proportion to exp( -(()) is the solution of the following optimization problem,
min qidensity
[Eg[f) + (og(()))
Hence we may interpret the Langevin algorithm as approximately solving an entropic regularized linear functional (i.c., free energy functional) on the probability space. This connection between

---

sampling and optimization (see Dalalyan (2017); Wibisono (2018); Durmus et al. (2019)) enables us employ the Langevin algorithm 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 aim to optimize is beyond the scope of Langevin algorithm due to the nonlinearity of loss {z, y) with respect to 2, the stationary distribution cannot be described as 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.
3 Proposed Method
We now propose the particle dual averaging method to approximately solve the problem (3) by optimizing two-layer neural network the mean field regime; we also introduce the mean field limit of the proposed method to explain the algorithmmic intuition and develop the convergence analysis.
3.1 Particle Dual Averaging
Our proposed particle dual averaging method (Algorithm 1) 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 more detailed explanation.
In the outer loop, the last iterate (())) of the previous inner loop is given. We compute 0. (Ir)), yt), which is a component of the Gibbs potential", and initialize set of parti- cles (1)) at (()) In Appendix 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 by (og(") which is sum of weighted average of 0. (z.), ye) (g)h)(( (6/5)) xx,) and the gradient of Ez-regularization (see Algorithm 1).
Algorithm 1 Particle Dual Averaging (PDA)
Input: data distribution D. initial density (d)), number of outer-iterations T, learning rates (m T
number of inner-iterations {T:)T=
Randomly draw i.i.d. initial parameters ((1) (O)) (r {1,2, M))
(()) + ( J,M
for t 1 to T do Randomly draw data (Iz, V1) from D (())
@(1) { |r=1
for k = 1 to Ti do
Run inexact noisy gradient descent for r E (1,2....,
V Veg (0(K))
(0(++2)043) Σt=1 (0g (2x), Mo )ph(O xx.) + 2h R!*))
Az((42)
g(*+1) 0 n (0!! + + 2t
(k) (i.i.d. Gaussian noise 5: IO)
end for
((++)) + (T++1) = (1(T++1)pp 1r=1
end for Randomly pick up t (2, ..., T 1} following the probability P|t] T(743)
2 and return har
Figure provides a pictorial illustration of Algorithm 1. Note that this procedure is slight modifica-
tion of the normal gradient descent algorithm: the first term of V o is similar to the gradient of the loss 0o, (hew (x),
dhe(he (æ), y) x)) where (R) (0. M Indeed, if
In Algorithm 1, the terms 0, (lhgr (æ.), y.) appear in inner loop: but note that these terms only need to be computed in outer loop because they are independent the inner loop iterates.
5

---

set the number of inner-iterations 7, = 1 and replace the direction V (o with the gradient of the Le-regularized loss, then PDA exactly reduces to the standard noisy gradient descent algorithm considered in Mci et al. (2018). Algorithm can be extended to the minibatch variant in the obvious manner, for efficient implementation in the empirical risk minimization setting see Appendix E. 1.
3.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 is the Langevin algorithm with M particles, which optimizes the potential function given by the weighted sum:
g (0) 20(1+2)(L+1) Σ (0.l(hgun
2
+ 1111013).
Due to the rapid convergence of Langevin algorithm outlined in Subsection 2.3, the particles of*++1) (r E [I....,,] MM) can be regarded as (approximate) samples from the Boltzmann distribution: exp (-g(0). Hence, the inner loop of PDA returns an A --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 + 0o 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 Dual Averaging (DA)
Input: data distribution D and initial density (1)
for t = 1 to T do
Randomly draw a data (It, y1) from D
(t) + 0_ (xz), (xx)), (e ()),(.;;;)) x4) + Aall
Obtain an approximation (1++2 of the density function q! ((+1) end for
Σ 22g
x exp
x2([+2]((-1
Randomly pick up t E (2,3, T 1} following the probability P[t]
2t TT]+3)
and return hyce
Algorithm 2 iteratively updates the density function qx
(++1)
E P2 which is a solution to the objective:
min
Eg
Σ sggn)]
X2 (t
2)(t + (og(())(()))
PEP:
where the function (()) 0.E(h (2t), Me (.., + All is the functional derivative of E(hy (xi;), vt) + A || with respect to at (1)) 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 approximating the optimal distribution q- by iteratively optimizing q=¹) which takes the form of a Boltzmann distribution. In the inner loop of the PDA algorithm, we obtain M (approximate) samples from gf via the Langevin algorithm. In other words, PDA can be viewed as a finite- maritcepattile approximation of DA indeed, the stationary distributions obtained in PDA converges to of (1+1) by taking M + co. 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; defer the partice approximation error analysis to Appendix C.

---

4 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.
4.1 Outer Loop Complexity
We first analyze the convergence rate the dual averaging (DA) method (Algorithm 2). Our analysis will be made under the following assumptions.
Assumption 1.
(A1) y c [-1,1] 1). ((2.y) is smooth convex function w.r. 2 and (0_((2,)) S f y, E.
(A2) [1(0,z)] - 1 and h(0, x) is smooth with respect to 0 for x E EX
(A3) (T( (1+1)) < 1/12.
Remark. (A2) is satisfied by smooth activation functions such sigmoid and tanh. Many loss functions including the squared loss and logistic loss satisfy (A1) under the boundedness assumptions y C 1-1. 1. 1] and [x)] can be relaxed to any value. (A3) specifies the precision of approximate solutions of sub-problems
1. Note that constants in (A1) and (A2) are defined for simplicity and
(7) to guarantee the global convergence of Algorithm 2, which we verify in our inner loop analysis.
We first introduce the following quantity for E P2.
<(q) def E.y [lg(a))
4 A2
kmm À1
exp
+ loo
A2
Observe that the expression consists of the negative entropy minus its lower bound for gh under Assumption (A1), (A2); in other words e(g!")) 2 0. We have the following convergence rate of DA
Theorem (Convergence of DA). Under Assumptions (A1), (A2), and (A3). for arbitrary q. Pa. iterates of the DA method (Algorithm 2) satisfies
2
T+1
Σt (E(((11))) C(g-))
T(T +3)
1=2
I 50 T (1+ [1011]) + dze(g-)
+ 11 cxp(8//22)) log (T 2)).
(
T
where the expectation [)) is taken with respect to the history of examples.
Theorem demonstrates the convergence rate of Algorithm 2 to the optimal value of the regularized
2 I++ (E) is the expectation of ((((g(())] (t (,...., T 17) specified in Algorithm If take p, AI, A2 as constants and use Ô to hide the logarithmic terms, we can deduce that after ( iterations, an e-accurate solution of the optimal distribution: L(q) < infgeP, 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).
objective (3) in expectation. Note that T(T+3)
according to the probability P[t] T(7+33
On the other hand, due to the exponential dependence on X, our convergence rate is not informative under weak regularization X2 + O. Such dependence follows from the classical LSI perturbation lemma (Holley and Stroock, 1987), which is likely unavoidable for Langevin- 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.
methods the
'In Appendix B we introduce a more general version of Theorem I that allows for inexact h,00 (z)), which simplifies the analysis of lnite-partiche discretization presented in Appendix C.
7

---

4.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):
Theorem 2 (Vempala and Wibisono (2019)). Consider a probability density q(0) x cxp(-f
satisfying the log-Sobolev inequality with constant a, and assume f is smooth and V f is L-Lipschitz i.e, [ - Vof (8') ll2 S Lil@ 8112- we run the Langevin algorithm (5) with learning rate 0 < 1) S 412 and let q (8) (0)d0 be a probability distribution that 61h) follows, then we have,
(L(g" < (xp(-on)))LL(()11 S "Sp''''''L'
Theorem 2 implies that a d-accurate solution in KL divergence can be obtained by the Langevin algorithm with 7) S _ min { 1,
& and co Tp
log 2KL(g" .tterttions.
-iterations.
(1+1) oc Since the optimal solution of sub-problem in DA (Algorithm 2) takes the forms of g= >> 2sg
xe(f+2)(e+1]/ we can verify the L.II and determine the constant for v- (0)d0 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 g= (1+1) for the inner loop complexity when Va log 9+ (1+11 is Lipschitz continuous, which motivates to introduce the following assumption.
exp
Assumption 2.
(A4) 0gh(,, æ²) is 1-Lipschitz continuous:
S 110-0'12.. Ve E 0, E Ω
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 Hölder continuity of Ogh(,,z) 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 8.+1 to be the desired accuracy of an approximate solution g(t+11 specified in (A3): dr+1 1/(t + 112,, we have the following guarantee for the inner loop.
Corollary (Inner Loop Complexity). Under (A1), (A2), and (A4), if we run the Langevin algo-
11 rithm with step size TT: O ((1+x1) exp(8/223) on (7), then an approximate solution satisfying
(L(g 1|ge (ltt)))) S 8.41 can be oblained within O Aa exp(8/2a)
log
2KL("")1955"") -iterations.
Aine
bet1
Moreover, (1 (1))) E (1, ................ 13) are uniformly bounded with respect to t as long as (2) 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 tatropolis-adjusste Langevin algorithm (Roberts and Tweedie, 1996; Dwivedi et al., 2018). which may improve the iteration complexity.
4.3 Total Complexity
Combining Theorem 1 and Corollary 1, we can now derive the total complexity of our proposed algorithm. For simplicity, we take p, Aι A2 constants and hide logarithmic terms in Ô and Ó. The following corollary establishes a Ō(E=a) total iteration complexity to obtain an e-accurate solution in expectation because T, @(t²) = Ō(c=2) for t $<5 ST
Corollary 2 (Total Complexity). Let e > 0 be an arbitrary desired accuracy and (1) be a Gaussian distribution. Under assumptions (A1), (A2), (A3), and (A4), if we run Algorithm f T Ô(e-)
\ A1dg iterations on the outer loop, and the Langevin algorithm with step size np e (
for T, = 0(n=)) iterations on the inner loop, then an e-accurate solution: L(q) S inl inf,Er, L(q) + of the objective (3) is achieved in expectation.
8

---

Quantitative convergence guarantee. To translate the above convergence rate result to the finite- particle PDA (Algorithm 1), we also characterize the fnite-partiche discretization enror in Appendix For the particle complexity analysis, we consider two versions of particle update: (i) the warm- start scheme described in Algorithm 1, in which θ is initialized at the last iterate O()) of the previous inner loop, and (ii) the resampling scheme, in which ( is initialized from the initial distribution q()) (0)d0 (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 Ō(c=2), together with the total iteration complexity of Ö(E-)), suffices obtain e-accurate solution to the objective (3) (see Appendix B and C for precise statement).
5 Experiments
5.1 Experiment Setup
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 Y/ f. (I;) + S,, 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:
M Σ
f(x)
(w, x by),
M
which we train to minimize the objective (3) using PDA. Note that a = 1 corresponds to the mean field regime (which are interested in), whereas setting a 1/2 leads the kernel (NTK) regime".
Synthetic student-teacher setting. For Figure ((a) we design synthetic experiments for both re- gression 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 multiple-index model with neurons: f- vm > a. ((wf, 2)), and the input is drawn from a unit Gaussian (0, For binary classifi- cation, we consider a simple two-dimensional dataset from sklearn. datasets.......... make_circles (Pedregosa al., 2011). I), in which the goal separate two groups of data on concentric circles (red and blue in Figure 2(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 T2 with f, and the step size 7) with 1/ I/Vi, where is the outer loop iteration; this heuristic is consistent with the required inner-loop accuracy Theorem and Proposition 2.
PDA hyperparameters.
10a
1.0
1.25
Taleing ab seccuue
0.9
10a
1.00
Train accuracy
Objective
accuracy
0.8
20-1
0.75
0.7
20-2
0.50
-1
"l clq-1
10-3
0(q"") - clq-)
00--1
-2
nogative
class PDA iteates
0.25
e17
0.5
0.6
0.00
107
10
150
Outer loop steps
Outer loop
(a) objective value
(b) parameter trajectory (classification).
(c) MNIST odd vs. even
(regression).
(classification).
Figure 2: (a) Iteration complexity PDA: O(T-1) rate on the outer loop agrees with Theorem 1. (b) Parameter trajectory of PDA: darker color (purple) indicates earlier in training, and vice versa. (c) odd vs. even
classification on MNIST; we report the training loss (red) as well as the train and test accuracy (blue and green).
We use the term kernel regime only to indicate the parameter scaling oc this does not necessarily imply that the NTK linearization accurate description of the trained model.

---

5.2 Empirical Findings
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 imput dimensionality p = 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 À 10-2; A2 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 parame- ters 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 good separator, its parameters should also distribute on 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 even digits from the MNIST dataset. We subsample 2500 training examples with binary labels, and learn a two-layer tanh network with width M = 2500. We use the resampled update of PDA optimize the cross entropy loss, with hyperparameters λ 10 X2 10 Observe that the algorithm achieves good generalization performance (green) and roughly maintains" the (T- 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 finiteddimensional 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 aligned with experimental results neural network optimization. Looking forward, we plan to investigate specific problem instances in which convergence rate can be obtained under vanishing regularization. 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 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 (PMMJR11928) DW was partially supported by NSERC and LG Electronics. TS was partially supported by JSPS KAKENHI (18H03201), Japan Digital Design and JST CREST.
References
Allen-Zhu, Z. and Li, Y. (2019). What resnet learn efficiently, going beyond kernels? In Advances in Neural Information Processing Systems, pages 9017-9028.
Allen-Zhup Z. and Li, Y. (2020). Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04411
Allen-Zhin Zhu, Z., Li, Y., and Song, Z. (2019). A convergence theory for deep learning via over- parameterization. In Proceedings of International Conference on Machine Learning 36, pages 242-252.
'Note that the estimated training objective (red) slightly deviates from the ideal 1/T rate; this may be due to inaccuracy in the entropy estimation, converereeeece of the algorithm (i.e.. overestimation of L(g.)).
10

---

Araújo, D., Oliveira, R. I., and Yukimura, D. (2019). A mean-field limit for certain deep neural networks. arXiv preprint
Bach, F. (2017). Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629-66) -81.
Bai, Y. and Lee, D. (2019). Beyond linearization: On quadratic and higher-order approximation of wide neural networks. arXiv preprint arXiv:
Bakry, D. and Emery, M. (1985). Diffusions hypercontraacivvs in sem, probab. xix Inm 1123.
Bengio, Y., Le Roux, N., Vincent, P., Delalleau, O., and Marcotte, P. (2006). Convex neural networks. Advances in neural information processing systems, 18:123.
Bou-Rabee, N. and Eberle, A. (2021). Mixing time guarantees for unadjusted hamiltonian monte carlo. arXiv e-prints, pages arXiv-2105.
Bou-Rabee, N. and Schuh, K. (2020). Convergence of unadjusted hamiltonian monte carlo for mean-field models. arXiv preprint arXiv:2009.08735.
Cao, Y. and Gu, Q. (2019). Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems, pages 10836-10846.
Chen, Z., Cao, Y., Gu, Q. and Zhang, T. (2020). A generalized neural tangent kernel analysis for two-layer neural networks. arXiv preprint 2oXi::2002.04026..
Cheng, X. and Bartlett, (2017). Convergence of langevin meme in kl-divergence. arXiv preprint 1705.00000.
Cheng, X., Chatterji, N. S., Bartlett, L., and Jordan, M. (2018). Underdamped langevin meme: A non-asymptotic analysis. In Conference on Learning Theory, pages 300-323. PMLR.
Chizat, L. (2019). Sparse optimization measures with over-parameterized gradient descent. arXiv preprint arXiv:1907.10300.
Chizat, L. (2021). Convergence rates of gradient methods for convex optimization in the space of measures, arXiv preprint 10X.v:2005.08368
Chizat, L. and Bach, F. (2018a). A note on lazy training in supervised differentiable programming. arXiv preprint arXiv:1812.0790
Chizat, and Bach, F. (2018b). On the global convergence of gradient descent for over peee-paraettriieed models using optimal transport. In Advances in Neural Information Processing Systems 31, pages 3040-3050.
Chu, C., Blanchet, J., and Glynn, P. (2019). Probability functional descent: A unifying perspective on gans, variational inference, and reinforcement learning. In Proceedings of International Conference on Machine Learning 36, pages 1213-1222.
Dai, B., He, N., Dai, H., and Song. L. (2016). Provable bayesian inference via particle mirror descent. In Proceedings of International Conference on Artificial Intelligence and Statistics 19, pages 985-994.
Dalalyan, A. S. (2014). Theoretical guarantees for approximate sampling from smooth and log- concave densities. arXiv preprint
Dalalyan, A. S. (2017). Further and stronger analogy between sampling and optimization: Langevin monte carlo and gradient descent. arXiv preprint 1704:.170
Daniely, A. and Malach, F. (2020). Learning parities with neural networks. arXiv preprint 2xXiv:2002.07400.
Ding, 7. and Li, Q. (2019). Ensemble kalman sampling: Mean-field limit and convergence analysis. arXiv preprint arXiv:1910.12923

---

Du, S. S., Zhai, X., Poczos, B., and Singh, A. (2019). Gradient descent provably optimizes over- parameterized neural networks. International Conference on Learning Representations 7.
Durmus, A., Majewski, S., and Miasojedow, B. (2019). Analysis of langevin monte carlo via convex optimization. Journal of Machine Learning Research, 20(73):1
Durmus, and Moulines, E. (2017). Nonasymptotic convergence analysis for the unadjusted langevin algorithm. The Annals of Applied Probability, 27(3):155--1587.
Dwivedi, R., Chen, Y., Wainwright, M. and Yu, B. (2018). Log-concave sampling: Metropolis- hastings algorithms are fast! In Conference on Learning Theory, pages 793-797. PMLR.
Eberle, A., Guillin, A., Zimmer, R., et al. (2019). Couplings and quantitative contraction rates for langevin dynamies. Annals of Probability, 4784):1922-2010.
Erdogdu, M. A. and Hosseinzadeh, R. (2020). On the convergence of langevin monte carlo: The interplay between tail growth and smoothness. arXiv preprint a3X77
Erdogdu, M. A., Mackey, L.,, and Shamir, O. (2018). Global non-convex optimization with discretized diffusions. In Advances in Neural Information Processing Systems, pages 9671-9680.
Garbuno-Inigo, A., Hoffmann, F., Li, W., and Stuart, A. M. (2020). Interacting langevin diffusions: Gradient structure and ensemble kalman sampler. SIAM Journal on Applied Dynamical Systems, 19(1):412-441.
Ghorbani, B., Mei, S., Misiakiewicz, T., and Montanari, A. (2019a). Limitations of lazy training of two-layers neural network. In Advances in Neural Information Processing Systems, pages 9111-9121.
Ghorbani, B., Mei, S., Misiakiewicz, T., and Montanari, A. (2019b). Linearized two-layers neural networks in high dimension. arXiv preprint arXiv: 1904.1211.
Ghorbani, B., Mei, S., Misiakiewicz, T., and Montanari, A. (2020). When do neural networks outperform kernel methods? arXiv preprint 20Xiv:2006.13409
Holley, R. and Stroock, D. (1987). Logarithmic sobolev inequalities and stochastic ising models. Journal of statistical physics, 46(--4):
Hsich, Y.PP., Liu, C., and Cevher, V. (2019). Finding mixed nash equilibria of generative adversarial networks. Proceedings of International Conference on Machine Learning 36, pages 2810-2819,
Hu, K., Ren, Z.. Siska, D., and Szpruch, L. (2019). Mean- field langevin dynamics and energy landscape of neural networks. arXiv preprint 19X5.0900500766
Imaizumi, M. and Fukumizu, K. (2020). Advantage of deep neural networks for estimating functions with singularity curves. arXiv preprint arXiv:2011.02256.
Jabir, J.FF, Šiška, D., and Szpruch, t.. (2019). Mean-field neural odes via relaxed optimal control. arXiv preprint arXiv: arXiv:1912.088 1912.05475.
Jacot, A., Gabriel, F., and Hongler, C. (2018). Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems 31, pages 8580-8589.
Javanmard, A., Mondelli, M., and Montanari, A. (2019). Analysis of a two-layer neural network via displacement convexity. arXiv preprint arXiv:1901.0137
Ji, Z. and Telgarsky. M. (2019). Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. arXiv preprint arXiv:1909.1222.
Jordan, R. and Kinderlehrer, D. (1996). 18. an extended variational. Partial differential equations and applications: collected papers honor of Carlo Pucci, 177:187.
Jordan, R., Kinderlehrer, D., and Otto, F. (1998). The variational formulation of the fokker-planck equation. SIAM journal mathematical analysis, 2911):1117
12

---

Kent, C., Blanchet, J., and Glynn, P. (2021). Frank-wolfe methods in probability space. arXiv preprint arXiv:210S.05352.
Kozachenko, L. and Leonenko, N. N. (1987). Sample estimate of the entropy of a random vector. Problemy Peredachi Informatsii, 23
Laurent, B. and Massart, P. (2000). Adaptive estimation of quadratic functional by model selection. The Annals of statistics, 28(5):1302-13
Li, X., Wu, Y., Mackey, L.,, and Erdogdu, M. A. (2019). Stochastic runge-kutta accelerates langevin monte carlo and beyond. In Advances in Neural Information Processing Systems, pages
Li, Y., Ma, T., and Zhang, H. R. (2020). Learning werrpaaameetrized two-layer neural networks beyond ntk. In Proceedings of Conference Learning Theory 33, pages 2613-2682.
Liu, Q. and Wang. D. (2016). Stein variational gradient descent: general purpose bayesian inference algorithm. In Advances in neural information processing systems, pages 2376..386.
Lu, J., Lu, Y., and Nolen, (2019). Scaling limit of the stein variational gradient descent: The mean field regime. SIAM Journal on Mathematical Analysis, 51)))6688 671.
Lu, Y., Ma, C.,, Lu, Y., Lu, J., and Ying, (2020). A mean-field analysis of deep resnet and beyond: Towards provable optimization via overparameterization from depth. arXiv preprint arXiv:2003.0553 05508.
Mattingly, J. C., Stuart, A. M., and Higham, D. J. (2002). Ergodicity for sdes and approximations: locally lipschitz vector fields and degenerate noise. Stochastic processes and their applications, 101(2):185-22
Mei, S., Montanari, A., and Nguyen, P.M. (2018). A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33),E765-E7777.
Menz, G. and Schlichting, A. (2014). Poincaré and logarithmic sobolev inequalities by decomposition of the energy landscape. The Annals of Probability, 42(5))1809-1188
Mohri, M., Rostamizadeh, A., and Talwalkar, A. (2012). Foundations of Machine Learning. The MIT Press.
Nesterov, Y. (2005). Smooth minimization of non-smooth functions. Mathematical programming, 107(55:
Nesterov, Y. (2009). Primal-dual subgradient methods for convex problems. Mathematical program- ming, 120(
Nguyen, P.-M. and Pham, H. T. (2020). A rigorous framework for the mean field limit of multilayer neural networks. arXiv preprint 200.:22001.11443:
Nitanda, A., Chinot, G., and Suzuki, T. (2019). Gradient descent can learn less over-parameterizeed two-layer neural networks on classification problems. arXiv preprint 4rXiv:1905.09870.
Nitanda, A. and Suzuki, T. (2017). Stochastic particle gradient descent for infinite ensembles. arXiv preprint arXiv: 171.00522
Nitanda, A. and Suzuki, T. (2020). Optimal rates for averaged stochastic gradient descent under neural tangent kernel regime. arXiv preprint arXhv:2006.12297,
Otto, and Villani, C. (2000). Generalization of an inequality by talagrand and links with the logarithmic sobolev inequality. Journal of Functional Analysis, 373(2):361-400.
Pedregosa, F,, Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M.,
Prettenhofer, P., Weiss, R., Dubourg, V., Vanderplas, J., Passos, A., Cournapeau, D., Brucher, M., Perrot, M., and Duchesnay, E. (2011). Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 2-2222-28800
13

---

Pham, H. T. and Nguyen, P.-M. (2021). Global convergence of three-layer neural networks in the mean field regime. In International Conference on Learning Representations.
Raginsky, M., Rakhlin, A., and Telgarsky, M. (2017). Non-convex learning via stochastic gradient langevin dynamics: nonasymptotic analysis. arXiv preprini arXiv:/702.0368
Roberts, G. O. and Tweedie, R. L. (1996). Exponential convergence of langevin distributions and their discrete approximations. Bernoulli, 2(4))341-363.
Rotskoff, G. M., Jelassi, S., Bruna, J., and Vanden-Fijjdent E. (2019). Global convergence neuron birth-death dynamics. In Proceedings of International Conference on Machine Learning 36, pages 9689-9698.
Rotskoff. G. M. and Vanden-Eijnden, E. (2018). Trainability and accuracy of neural networks: An
interacting particle system approach. arXiv preprint /
Schmidt-Hieber, J. (2020). Nonparametric regression using deep neural networks with relu activation function. The Annals of Statisties, 48(4):1875-1897.
Shalev-Shwartz, S. and Ben-David, S. (2014). Understanding machine learning: From theory to algorithms. Cambridge university press.
Sirignano, J. and Spiliopoulos, K. (2020). Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications, 13203)2120--152.
Suzuki, (2018). Adaptivity deep relu network for learning in besov and mixed smooth besov spaces: optimal rate and curse of dimensionality, arXiv preprint arXiv: /810.008100
Suzuki, T. (2020). Generalization bound of globally optimal non-convex neural network train- ing: Transportation map estimation by infinite dimensional langevin dynamics. arXiv preprint arXiv:2007.05824.
Suzuki, T. and Nitanda, A. (2019). Deep learning is adaptive intrinsic dimensionality of model smoothness in anisotropic besov space. arXiv preprint 19X.:1910.
Vempala, S. and Wibisono, A. (2019). Rapid convergence of the unadjusted langevin algorithm:
Isoperimetry suffices. In Advances Neural Information Processing Systems, pages 8094-8106
Wei, C., Lee, J. D., Liu, Q., and Ma, T. (2019). Regularization matters: Generalization and optimization of neural nets vs their induced kernel. In Advances in Neural Information Processing Systems, pages 9712-9724.
Wibisono, A. (2018). Sampling as optimization in the space of measures: The langevin dynamics as composite optimization problem. In Proceedings of Conference on Learning Theory 31, pages 2093-3027.
Xiao, L. (2009). Dual averaging method for regularized stochastic learning and online optimization. In Advances in Neural Information Processing Systems 22, pages 2116-2124.
Xu, P., Chen, J., Zou, D., and Gu, Q. (2018). Global convergence of langevin dynamics based algorithms for nonconvex optimization. In Advances in Neural Information Processing Systems, pages 3122-3133. 3133.
Yang. G. and Hu, E. J. (2020). Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011.14522.
Ychudai, G. and Shamir, O. (2019). On the power and limitations of random features for understanding neural networks. In Advances in Newral Information Processing Systems, pages 6598 -6088.
Ying, L. (2020). Mirror descent algorithms for minimizing interacting free energy. Journal of Scientific Computing, 84(3):1-14.
Zou, D., Cao, Y., Zhou, D., and Gu, Q. (2020). Gradient descent optimizes over- over- wer-parameterized deep relu networks. Machine Learning. 19943):467-492.
14


Reviews:
[{"review_id": "rjzlltwlDKC", "reviewer": "Reviewer_fmps", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "STRENGTHS:\nFirst of all the paper is very well-written. The motivation and the algorithm are clearly presented and the theoretical results are given with precise assumptions. \n\nThe idea of using dual averaging for probabilistic optimization is original and really promising. One of the key remark that minimizers of linearized functionals on probability measures with Kullback-Leibler regularization are given by Gibbs measures might find applications outside of the context of overparameterized neural networks.\n\nThe authors use recent tools in optimization/sampling (convergence in Kullback-Leibler divergence of the Langevin algorithm under log-Sobolev condition with quantitative rates, convergence of the dual averaging algorithm) in order to derive quantitative results. These quantitative results are really promising. \n\nWEAKNESSES:\nOne of the main weakness of the paper is its experiments which I found to be a bit too limited even for a theoretical paper. The authors could have dealt with real datasets for classification such as MNIST or CIFAR-10 in order to assess the scalability of their method with respect to the dimension and hence its applicability.\n\nSimilarly, even if some comparison with a classical training of neural networks (with SGD) are reported in the appendix, I think further comparisons should be conducted to assess the efficiency of the proposed methodology (for instance in Figure 4 in the appendix, PDA seems to perform only marginally better than SGD which is simpler to train). Also instead of SGD the author might consider sampling with a Langevin algorithm (i.e. changing the noise scaling). Doing so incorporates an entropic term and will allow for a fairer comparison with PDA in my opinion.\n\nFrom a theoretical point of view I think that the authors should discuss in greater details the relationship between their algorithm and other works. In particular, how does the proposed approach relates to the ones of [1,2] which prove theoretical results for either SGD or the Langevin dynamics and not the PDA algorithm.\n\n[1] Hu, Ren, Siska, Szpruch - Mean-field Langevin dynamics and energy landscape of neural networks\n[2] Javanmard, Mondelli, Montanari - Analysis of a two-layer neural network via displacement convexity\n\nCOMMENTS:\n-The original dual averaging algorithm allows flexibility in the primal and dual discretization steps, see (2.14) in [3]. In this work the authors have chosen $\\beta_k = O(k^2)$ and $\\lambda_k = O(k)$ if I am not mistaken. This choice of parameters is not discussed in the paper.\n-Related to this remark, one of the limitation of the algorithm is that each inner-loop gets more and more expensive. Is this due to the way the parameters of the dual averaging are set? Is there any way to mitigate this issue?\n\nMINOR COMMENTS:\n-I am not sure to understand the role of $n$ in Proposition D. I assume that $x_i$ are the generated the test data points?\n\n[3] Nesterov - Primal-dual subgradient methods for convex problems", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Societal impact is irrelevant for this very theoretical work. The limitations of the algorithm could have been better highlighted and addressed (increasing cost of the inner loop, marginal performance over SGD), see my previous comments.", "review_text": "", "extra_scores": {"ethical_concerns": "None ", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "20"}, "extra_text": {}}, {"review_id": "n_vh4i-KRGF", "reviewer": "Reviewer_QjGu", "summary": "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. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper  develops optimization algorithms called PDA for neural networks in the mean field regime with \naccurate quantitative guarantees as  the kernel regime enjoys,  which is an important issue in understanding the \noptimization error in deep learnling.  Althrouth the ideas and techniques used in this paper are not new, the authors \nestablish quantitative global convergence rate of PDA in minimizing an KL-regularized objective  for the two layer neural network opitmization.   Overall,  this paper gives a nice contribution in understanding optimization in deep learning. \n\nOne question: Does the constant hidden in $\\tilde{O}(\\epsilon^3)$ depend on the ambient dimension? If so, what is the quantitive \nrelation between them? ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "cuzMFqS6EbF", "reviewer": "Reviewer_gSde", "summary": "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$.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper makes a serious contribution to the growing literature on the global convergence of two layer neural networks, and is one of the first to provide a quantitative error bound, albeit under a somewhat more restrictive setting (requires entropy regularization with a large constant for the bound to be meaningful).\nThe paper is well-written and an effort is made to make it accessible to a broader audience, and this effort is commendable.\n\nPerhaps because of this, most of the technical contribution is unfortunately pushed to the appendix, and there is no explicit error bound for PDA (the closest is Remark line 782 in the appendix). In other words, the main results (Theorem 1 and Corollary 2) concern an idealized mean field version (Algorithm 2) that is not implementable. I think it would be much better to write an explicit bound for PDA in the main paper.\n\nOn the exponential dependence on 1/\\lambda_2: thank you for acknowledging this limitation and pointing to precisely why it arises. In the comment line 243, it is claimed that this dependence \"is unfortunately unavoidable\", is this alluding to a lower bound? If so, please include a reference, otherwise, please rephrase to make it less factual (for example \"we believe that...\") and provide a reasoning why this might be.\n\nAn additional point to clarify: line 665, it is mentioned that when considering the warm-start scheme, $h_x$ is set to $h_q$ ... without tolerance ($\\epsilon = 0$). Why is it fine to assume $\\epsilon = 0$ in this case? Even in this case one only gets a particle approximation of  q. Please clarify.\n\nAdditional suggestions:\n- Consider adding a remark that the Gibbs distributions (such as $q_*^{(t)}$) are well defined thanks to L2 regularization and assumption (A1). This will clarify some of the importance of regularization.\n- Appendix A: clarify what definition of gradient you are using, since $P_2$ is not a vector space.\n- Overloaded notation: p denotes a density and the dimension of the parameter space (e.g. Definition A)\n- Line 647: typo.\n- Line 680: optimal coupling is not defined.\n- The second and third experiments in the main paper don't seem to relate to the main narrative. They may be moved to the appendix to make additional space for an explicit bound on PDA.\n- Experiment G.1: when comparing to SGD, it seems appropriate to also use entropy regularization (i.e. compare to Langevin gradient descent).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "15"}, "extra_text": {}}, {"review_id": "5VEYJsnlO3", "reviewer": "Reviewer_kvcd", "summary": "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. ", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality:  The idea of using a dual averaging algorithm (mainly used in online learning literature) is interesting. The convergence of the approximating sequence (outer loop) is established in Theorem 1, and this seems like the main contribution of this work.  Authors admit the proof idea is the same as in Nesterov 2009 and Xiao 2009, but I'm not familiar with these works.  What's puzzling is Assumption A3 that does not seem to be compatible with Theorem 2 in terms of rate. Furthermore, the authors never verify (at least not in Example 2 and Lemma B as they refer to) that each measure in the sequence satisfies log-Sobolev inequality (with what constant) and do not study the impact of the regularisation \\lambda_1 on the rate of convergence.  Hu et al. already established an exponential rate of convergence in the case when \\lambda_1 is sufficiently big.  \n\nIn line 38 authors claim that continuous-time results cannot be easily translated to a discrete setting, but this is not true. See for example:\nJabir, J. F., Siska, D., & Szpruch, L. (2019). Mean-field neural odes via relaxed optimal control. arXiv preprint arXiv:1912.05475\nfor an example for such analysis in a fairly general setting. There is large body of research on uniform in time propagation of chaos and discretisation error for mean-field Langevin dynamics e.g:\n- Majka, Mateusz B., Aleksandar Mijatović, and Lukasz Szpruch. \"Nonasymptotic bounds for sampling algorithms without log-concavity.\" The Annals of Applied Probability 30.4 (2020): 1534-1581\n- Delarue, François, and Alvin Tse. \"Uniform in time weak propagation of chaos on the torus.\" arXiv preprint arXiv:2104.14973 (2021).\n\nClarity:  The presentation of the paper could improve in places. Some suggestions:\nin section 2.1 the authors should stick to 1-hidden layer example. The ensemble of networks is confusing and not so interesting.\nI struggled to gain intuition about what the dual averaging method is from section 3.1 and needed to go back to the literature to see what is going on.  Maybe authors could expand section 2.3 and explain dual averaging there. \nAssumption A3 is in fact, a result that needs to be established. The authors should explain how this is done in Theorem 2\n\nSignificance:  Overall, this is an interesting research direction, but the presentation would need to improve, and a more thorough comparison with the literature needs to be presented i.e can authors can get away with less regularisation than Hu et al. and still obtain the rate of convergence? ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "n/a", "review_text": "", "extra_scores": {"ethical_concerns": "n/a", "needs_ethics_review": "No", "ethics_review_area": ["I don’t know"], "time_spent_reviewing": "3h"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors introduce a novel Particle Dual Averaging technique to minimize to minimize a non-linear functional of probability distributions. In terms of applications, they focus on the training of two-layer neural networks. In this context, they present quantitative convergence results for an entropy-regularized loss function. They also discuss the finite particle algorithm and present generalization bounds. 

There was a short but informative discussion between reviewers. All the reviewers agree that the paper presents some original, interesting and solid theoretical contributions. However, it was felt by some reviewers that the experimental section was fairly weak and I encourage the authors to improve it before publication.

I recommend acceptance of the paper.


Decision: Accept (Poster)