Mean-Field Langevin Dynamics and Energy Landscape of Neural Networks

Abstract

We present a probabilistic analysis of the long-time behaviour of the nonlocal, diffusive equations with a gradient flow structure in 2-Wasserstein metric, namely, the Mean-Field Langevin Dynamics (MFLD). Our work is motivated by a desire to provide a theoretical underpinning for the convergence of stochastic gradient type algorithms widely used for non-convex learning tasks such as training of deep neural networks. The key insight is that the 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 show that the flow of marginal laws induced by the MFLD converges to the stationary distribution which is exactly the minimiser of the energy functional. We show that this convergence is exponential under conditions that are satisfied for highly regularised learning tasks. At the heart of our analysis is a pathwise perspective on Otto calculus used in gradient flow literature which is of independent interest. Our proof of convergence to stationary probability measure is novel and it relies on a generalisation of LaSalle's invariance principle. Importantly we do not assume that interaction potential of MFLD is of convolution type nor that has any particular symmetric structure. This is critical for applications. 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.

Convex Analysis of the Mean Field Langevin Dynamics

“…• By analyzing the proximal Gibbs distribution p q , we establish linear convergence in continuous time with respect to the KL-regularized objective. This convergence result holds for any regularization parameters, in contrast to existing analyses (e.g., Hu et al (2019)) that require strong regularization.…”

“…This fact was shown in Hu et al (2019). Hence, we may interpret the divergence between probability density functions q and p q as an optimization gap.…”

“…Quantitative convergence rate usually requires additional conditions, such as structural assumptions on the learning problem (Javanmard et al, 2019;Chizat, 2021b;Akiyama and Suzuki, 2021), or modification of the dynamics (Rotskoff et al, 2019;Wei et al, 2019). Noticeably, Mei et al (2018); Hu et al (2019); Chen et al (2020) considered the optimization of the KL-regularized objective, which leads to the mean field Langevin dynamics.…”

“…Recent works have studied the convergence rate of the mean field Langevin dynamics, including its underdamped (kinetic) version. However, most existing analyses either require sufficiently strong regularization (Hu et al, 2019;Jabir et al, 2019), or build upon involved mathematical tools (Kazeykina et al, 2020;Guillin et al, 2021). Our goal is to provide a simpler convergence proof that covers general and more practical machine learning settings, with a focus on neural network optimization in the mean field regime.…”

Relativistic Langevin Equation derived from a particle-bath Lagrangian

“…The Langevin equation is the cornerstone of modern nonequilibrium statistical mechanics [1,2], with profound implications for numerical computations [3,4]. It is the basis upon which one builds the diffusion equations [5][6][7][8], fluctuation-dissipation theorems [1,9], Brownian motion with state-dependent diffusion coefficient and its path-integral formulation [10], and instrumental in areas overlapping with other fundamental theories of physics, such as the question of quantum decoherence [11,12].…”