Particle Stochastic Dual Coordinate Ascent: Exponential convergent algorithm for mean field neural network optimization

1 INTRODUCTION . A major challenge in developing an optimization theory for neural network is to handle the nonconvexity of the loss landscape . Recent works showed that overparameterization is one of the key properties that makes global optimization possible . In particular , the effectiveness of overparameterization has been extensively investigated in the theory of Neural Tangent Kernel ( NTK ) ( Jacot et al. , 2018 ; Du et al. , 2019 ; Allen-Zhu et al. , 2019 ) and the mean field analysis ( Nitanda & Suzuki , 2017 ; Mei et al. , 2018 ; Chizat & Bach , 2018 ; Sirignano & Spiliopoulos , 2020 ) . The NTK analysis considers a linear approximation of neural networks and casts the gradient descent dynamics to that in a corresponding reproducing kernel Hilbert space ( Jacot et al. , 2018 ; Du et al. , 2019 ; Allen-Zhu et al. , 2019 ) . Since the whole argument can be carried out essentially in a linear space , we can derive both global convergence and generalization guarantees . However , one of the biggest drawbacks of the NTK approach is that the fixed kernel fails to capture the feature training aspect of deep learning . Indeed , this feature learning ability of neural network is an important ingredient that separates deep learning from shallow models such as kernel methods ( Suzuki , 2018 ; Yehudai & Shamir , 2019 ; Ghorbani et al. , 2019 ) . In contrast , the mean field analysis considers a smaller scale of parameters than the NTK analysis , which allows the parameters to travel away from the initialization and learn features . However , the optimization problem is no longer convex and the analysis requires much more involved mathematical tools . A typical approach is to formulate the optimization of neural networks as a partial differential equation ( PDE ) of the distribution of parameters , for which global convergence can be established under certain conditions ( Nitanda & Suzuki , 2017 ; Mei et al. , 2018 ; Chizat & Bach , 2018 ; Rotskoff & Vanden-Eijnden , 2018 ; Sirignano & Spiliopoulos , 2020 ) . A benefit of the PDE formulation is that the objective function becomes convex with respect to the probability measure of the parameters . However , the dynamics is essentially infinite dimensional , and thus most convergence results are shown in the continuous limit ( infinite-width limit ) . Moreover , time discretization of the corresponding PDE is also challenging , so most existing works dealt with a continuous time dynamics . Indeed , Nitanda & Suzuki ( 2017 ) ; Chizat & Bach ( 2018 ) ; Rotskoff & Vanden-Eijnden ( 2018 ) proved convergence to a global optimal solution for essentially infinite width networks under continuous time settings . On the other hand , global optimality may be shown under less restrictive settings if we consider the noisy gradient descent algorithm ; the corresponding stochastic dynamics can be formulated as McKean–Vlasov dynamics , and its convergence to a global solution has been studied in Mei et al . ( 2018 ) ; Hu et al . ( 2019 ) . However , most of these works establish convergence result for the discrete-time finite-width method by bounding the difference from its continuous-time and infinite-width limit counterpart ; this usually incurs a large discretization error , which presents a limitation in its efficiency . Hence , it is important to develop a practical algorithm with optimization guarantees for finite width network and discrete time setting . Recently , Nitanda et al . ( 2021 ) proposed a completely time-space discretized algorithm termed Particle Dual Averaging ( PDA ) that attains polynomial time convergence guarantee to the ( regularized ) global optimal solution . This method combines the dual averaging technique ( Nesterov , 2005 ; 2009 ; Xiao , 2009 ) and the gradient Langevin dynamics ( Vempala & Wibisono , 2019 ) . More specifically , the algorithm realized time discretization by extending the dual averaging technique for finite dimensional convex optimization to the infinite dimensional setting , and realized the space discretization by sampling finite number of particles via the gradient Langevin dynamics . Eventually , the method achieves O ( 1/t ) convergence with respect to the outer-loop iteration t. Although O ( 1/t ) is optimal in an online optimization setting , there is still room for improvement for a finite sample size setting . Motivated by this observation , we propose a method that improves the convergence rate from the polynomial order ( O ( 1/t ) ) to an exponential order ( O ( exp ( −Ct ) ) ) . Our contribution . We propose Particle-SDCA , a novel optimization algorithm for two-layer neural networks in the mean field regime that achieves linear convergence ( w.r.t . the outer loop steps ) for regularized empirical risk minimization1 . We adopt the stochastic dual coordinate ascent ( SDCA ) technique ( Shalev-Shwartz & Zhang , 2013a ; Takáč et al. , 2013 ; Shalev-Shwartz & Zhang , 2013b ) to achieve the linear convergence rate and much better dependency on the sample size n. Importantly , the integration with respect to the probability measure on the parameters is approximated by a finite number of particles . In contrast to existing methods , in which particles are sampled at each iteration , we employ a novel update rule of the weights on the particles so that the sampling procedure is only performed once in ñ-iterations , which provides significant computational speedup . Furthermore , unlike existing analyses concerning the time-space discretized version of McKean–Vlasov dynamics ( Mei et al. , 2018 ; Bou-Rabee & Schuh , 2020 ; Bou-Rabee & Eberle , 2021 ) , we do not couple the discrete time update with the continuous time counterpart by taking the small step size limit . Instead , we directly analyze the convergence of the discrete time update ; hence , our method converges to the optimal solution with a fixed ( non-vanishing ) step size . Our contribution can be summarized as follows ( for comparison with existing works see Table 1 ) : • We propose a new optimization method for mean field representation of neural networks that is efficient in the finite sample setting . By utilizing the SDCA technique , the algorithm achieves exponential convergence with respect to the outer iteration . Moreover , the global convergence is guaranteed without letting the step size converge to 0 . • The integral required in each update is approximated by an average over finite number of particles . The number of particles can be of polynomial order with respect to a required precision , and the computational cost for re-sampling can be of poly-log order . • Thanks to a novel re-weighting scheme for particles , we only resample particles once in ñ updates , where ñ can be as large as the total sample size n ; this significantly reduces the computation cost . 1We note that our goal is to propose a more efficient algorithm for optimizing mean field neural network with faster convergence guarantee , rather than to gain a better understanding on standard neural network training . 2 PRELIMINARIES . We first introduce the problem setting and some notations used in the following sections . We consider the empirical risk minimization problem in the supervised learning setting with a twolayer neural networks , where the input and output spaces are denoted by X ⊂ Rd and Y ⊂ R , respectively . Let hθ : X → Y be one neuron with the parameters θ ∈ Rd̃ . For example , we may set hθ ( x ) = tanh ( r ) σ ( w > x ) for θ = ( r , w ) ∈ R×Rd and an activation function σ ( the tanh operation is merely to ensure boundedness of the output ) . In the mean field regime , we consider neural networks represented as an average of neurons : fΘ ( x ) = 1M ∑M m=1 hθm ( x ) where Θ = ( θm ) M m=1 ⊂ Rd̃ is a set of parameters . The continuous limit in the mean field regime is obtained by taking M →∞ , and by an analogy to the law of large numbers , the function fΘ converges to the following integral representation of a two-layer neural network : fp ( x ) = ∫ hθ ( x ) p ( θ ) dθ , where p : Rd̃ → R is a probability density function on Rd̃ representing the weight of parameters . Here , we denote by P the set of probability density functions on Rd̃ . As in the typical mean field analysis ( Nitanda & Suzuki , 2017 ; Mei et al. , 2018 ; Chizat & Bach , 2018 ) , we aim to optimize the density function p ∈ P so that the neural network fp accurately predicts the output y ∈ Y from the input x ∈ X . To define the empirical risk and predictive risk , we let ` ( z , y ) : Y × Y → R be a convex loss function , such as the squared loss ` ( z , y ) = ( z − y ) 2/2 for regression , and the logistic loss ` ( z , y ) = log ( 1 + exp ( −yz ) ) for classification . For each ( xi , yi ) , we use the notation hi ( θ ) and ` i ( f ( xi ) ) to indicate hθ ( xi ) and ` ( f ( xi ) , yi ) respectively . To estimate the density function p ∈ P , we minimize the regularized empirical risk defined by min p∈P 1 n n∑ i=1 ` i ( fp ( xi ) ) + λKL ( p||N ( 0 , σ2I ) ) ( 1 ) where λ > 0 is a regularization parameter and KL ( p , N ( 0 , σ2I ) ) represents the KL divergence from the normal distributionN ( 0 , σ2I ) with mean 0 and covariance σ2I to the distribution with the density function p. This type of regularization naturally arises in training neural networks by the gradient Langevin dynamics ( Hu et al. , 2019 ; Chen et al. , 2020b ) . In our framework , we consider directly optimizing the empirical risk with KL regularization , instead of running gradient Langevin dynamics ( i.e. , noisy gradient descent ) on the unregularized objective . By decomposing the KL divergence in Eq . ( 1 ) , we obtain the following equivalent representation of the objective : min p∈P P ( p ) : = 1 n n∑ i=1 ` i ( fp ( xi ) ) + λ1 ∫ p ( θ ) ‖θ‖2dθ + λ2 ∫ p ( θ ) log ( p ( θ ) ) dθ , ( 2 ) where λ1 , λ2 > 0 are the regularization parameters . Importantly , the objective is convex with respect to the density p ∈ P although it is non-convex with respect to each parameter θ . We make full use of this convexity in the following analysis . The main difficulty to optimize ( 2 ) stems from the following two factors : ( i ) there is an integral with respect to p that should be approximated by a computationally tractable scheme , and ( ii ) there is a summation over n data points , which could lead to costly gradient evaluations . To overcome the former difficulty , we sample a finite number of particles from the density p. As for the later , we adopt the stochastic dual coordinate ascent ( SDCA ) technique , which we explain in the next section .

Summary: The problem of training an infinitely wide neural network with one hidden layer can be written as an optimization problem in the space of probability distributions, where each atom in the distribution corresponds to a hidden neuron. The paper proposes to solve instead the dual problem (of an entropy regularized variant), which in the case of finite data is a finite-dimensional optimization problem.  The main contribution is the conception and analysis of an efficient dual coordinate ascent scheme to solve the dual, which in theory and practice requires fewer iterations to find an epsilon-accurate solution than existing methods.