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 paper considers optimization of MF two-layer neural nets (the infinite-width version with entropic regularization). Specifically the paper proposes to do dual coordinate ascent, which allows for exponential convergence rate, together with a particle approximation scheme. The paper shows that in the discrete-time setting, an improved convergence rate is obtained (w.r.t. desired error $\epsilon_P$, the number of data points $n$, the regularization $\lambda_2$).