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

1 Introduction . Gradient-based optimization can achieve vanishing training error on neural networks , despite the apparent non-convex landscape . Among various works that explains the global convergence , one common ingredient is to utilize overparameterization to translate the training dynamics into function spaces , and then exploit the convexity of the loss function with respect to the function . Such endeavors usually consider models in one of the two categories : the mean field regime or the kernel regime . On one hand , analysis in the kernel ( lazy ) regime connects gradient descent on wide neural network to kernel regression with respect to the neural tangent kernel ( Jacot et al. , 2018 ) , which leads to global convergence at linear rate ( Du et al. , 2019 ; Allen-Zhu et al. , 2019 ; Zou et al. , 2020 ) . However , key to the analysis is the linearization of the training dynamics , which requires appropriate scaling of the model such that distance traveled by the parameters vanishes ( Chizat and Bach , 2018a ) . Such regime thus fails to explain the feature learning of neural networks ( Yang and Hu , 2020 ) , which is believed to be an important advantage of deep learning ; indeed , it has been shown that deep learning can outperform kernel models due to this adaptivity ( Suzuki , 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 . 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 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 . 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 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 . 1.3 Notations . 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 . 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 ⇢ 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 . 2.1 Neural Network and Mean Field Limit . 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 .

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.