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

TL;DR

This paper introduces a novel method to simulate Fokker-Planck equations using mean field control and score-based normalizing flows, enabling efficient density evolution modeling for stochastic systems.

Contribution

It formulates FP equations as a mean field control problem and leverages score-based normalizing flows for efficient computation of score functions along trajectories.

Findings

Effective simulation of FP equations for various stochastic systems.

Convergence demonstrated for Ornstein-Uhlenbeck processes.

Numerical validation on Langevin and chaotic systems.

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.