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

Abstract

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

NeurIPS 2021

References in corpus (10)

• A mean-field limit for certain deep neural networks
• A Generalized Neural Tangent Kernel Analysis for Two-layer Neural Networks
• Stochastic Particle Gradient Descent for Infinite Ensembles
• A Mean-field Analysis of Deep ResNet and Beyond: Towards Provable Optimization Via Overparameterization From Depth
• Mixing Time Guarantees for Unadjusted Hamiltonian Monte Carlo
• Convergence of unadjusted Hamiltonian Monte Carlo for mean-field models
• On the Convergence of Langevin Monte Carlo: The Interplay between Tail Growth and Smoothness
• Learning Parities with Neural Networks
• Learning Over-Parametrized Two-Layer ReLU Neural Networks beyond NTK
• Global Convergence of Three-layer Neural Networks in the Mean Field Regime