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

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 relies on the derivatives of the parameterized velocity fields. Numerical results, including Langevin dynamics, underdamped Langevin dynamics, chaotic systems, and high-dimensional interacting particle systems validate the effectiveness and scalability of our proposed algorithm. A convergence analysis is conducted for our algorithm on the FP equation of Ornstein--Uhlenbeck processes.