Simulating Fokker-Planck equations via mean field control of score-based normalizing flows

Abstract

The Fokker-Planck (FP) equation governs the evolution of densities for stochastic dynamics of physical systems, such as the Langevin dynamics and the Lorenz system. This work simulates FP equations through a mean field control (MFC) problem. We first formulate the FP equation as a continuity equation, where the velocity field consists of the drift function and the score function, i.e., the gradient of the logarithm of the density function. Next, we design a MFC problem that matches the velocity fields in a continuity equation with the ones in the FP equation. The score functions along deterministic trajectories are computed efficiently through the score-based normalizing flow, which only rely on the derivatives of the parameterized velocity fields. A convergence analysis is conducted for our algorithm on the FP equation of Ornstein-Uhlenbeck processes. Numerical results, including Langevin dynamics, underdamped Langevin dynamics, and various chaotic systems, validate the effectiveness of our proposed algorithms.