Score-based Neural Ordinary Differential Equations for Computing Mean Field Control Problems

Classical neural ordinary differential equations (ODEs) are powerful tools for approximating the log-density functions in high-dimensional spaces along trajectories, where neural networks parameterize the velocity fields. This paper proposes a system of neural differential equations representing first-and second-order score functions along trajectories based on deep neural networks. We reformulate the mean field control (MFC) problem with individual noises into an unconstrained optimization problem framed by the proposed neural ODE system. Additionally, we introduce a novel regularization term to enforce characteristics of viscous Hamilton-Jacobi-Bellman (HJB) equations to be satisfied based on the evolution of the second-order score function. Examples include regularized Wasserstein proximal operators (RWPOs), probability flow matching of Fokker-Planck (FP) equations, and linear quadratic (LQ) MFC problems, which demonstrate the effectiveness and accuracy of the proposed method. Score functions have been widely used in modern machine learning algorithms, particularly generative models through time-reversible diffusion (Song et al., 2021). The score function can be viewed as a deterministic representation of diffusion in stochastic trajectories (Carrillo et al., 2019). These properties have inspired algorithms for simulating stochastic trajectories or sampling problems that converge to target distributions (Wang et al., 2022; Lu et al., 2024). Typical applications include modeling the time evolution of probability densities for stochastic dynamics and solving control problems constrained by such dynamics. While score functions provide powerful tools for modeling stochastic trajectories, their computations are often inefficient, especially in high-dimensional spaces. Classical methods, such as kernel density estimation (KDE) (Chen, 2017), tend to perform poorly in such settings due to the curse of dimensionality (Terrell & Scott, 1992). Recently, neural ODEs have emerged as efficient ways of estimating densities. In particular, one uses neural networks to parameterize the velocity fields and then approximates the logarithm of density function along trajectories. The time discretizations of neural ODEs can be viewed as normalization flows in generative models.