Mean-Field Langevin Dynamics and Energy Landscape of Neural Networks

Abstract

Our work is motivated by a desire to study the theoretical underpinning for the convergence of stochastic gradient type algorithms widely used for non-convex learning tasks such as training of neural networks. The key insight, already observed in the works of Mei, Montanari and Nguyen (2018), Chizat and Bach (2018) as well as Rotskoff and Vanden-Eijnden (2018), is that a certain class of the finite-dimensional non-convex problems becomes convex when lifted to infinite-dimensional space of measures. We leverage this observation and show that the corresponding energy functional defined on the space of probability measures has a unique minimiser which can be characterised by a first-order condition using the notion of linear functional derivative. Next, we study the corresponding gradient flow structure in 2-Wasserstein metric, which we call Mean-Field Langevin Dynamics (MFLD), and show that the flow of marginal laws induced by the gradient flow converges to a stationary distribution, which is exactly the minimiser of the energy functional. We observe that this convergence is exponential under conditions that are satisfied for highly regularised learning tasks. Our proof of convergence to stationary probability measure is novel and it relies on a generalisation of LaSalle's invariance principle combined with HWI inequality. Importantly, we assume neither that interaction potential of MFLD is of convolution type nor that it has any particular symmetric structure. Furthermore, we allow for the general convex objective function, unlike, most papers in the literature that focus on quadratic loss. Finally, we show that the error between finite-dimensional optimisation problem and its infinite-dimensional limit is of order one over the number of parameters.

arXiv comment31 pages

References in corpus (9)

• Understanding deep learning requires rethinking generalization
• A Mean Field View of the Landscape of Two-Layers Neural Networks
• Machine learning approximation algorithms for high-dimensional fully nonlinear partial differential equations and second-order backward stochastic differential equations
• A proof that artificial neural networks overcome the curse of dimensionality in the numerical approximation of Black-Scholes partial differential equations
• Uniform convergence to equilibrium for granular media
• Does data interpolation contradict statistical optimality?
• Weak quantitative propagation of chaos via differential calculus on the space of measures
• Trajectorial Otto calculus
• Exponential ergodicity for SDEs and McKean-Vlasov processes with Lévy noise

Cited by in corpus (15)

• Game-Theoretic Multiagent Reinforcement Learning
• Towards a Mathematical Understanding of Neural Network-Based Machine Learning: what we know and what we don't
• Mathematical Models of Overparameterized Neural Networks
• A Mean-field Analysis of Deep ResNet and Beyond: Towards Provable Optimization Via Overparameterization From Depth
• Can Shallow Neural Networks Beat the Curse of Dimensionality? A mean field training perspective
• Mean-Field Neural ODEs via Relaxed Optimal Control
• On the Convergence of Gradient Descent Training for Two-layer ReLU-networks in the Mean Field Regime
• Mean-field Langevin System, Optimal Control and Deep Neural Networks
• Stochastic gradient descent and fast relaxation to thermodynamic equilibrium: a stochastic control approach
• Ergodicity of the underdamped mean-field Langevin dynamics
• Particle Dual Averaging: Optimization of Mean Field Neural Networks with Global Convergence Rate Analysis
• On the Empirical Neural Tangent Kernel of Standard Finite-Width Convolutional Neural Network Architectures
• On the exponential ergodicity of the McKean-Vlasov SDE depending on a polynomial interaction
• Non-asymptotic estimation of risk measures using stochastic gradient Langevin dynamics
• A trajectorial approach to relative entropy dissipation of McKean$-$Vlasov diffusions: gradient flows and HWBI inequalities