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

TL;DR

This paper introduces a neural ODE framework for mean field control problems, leveraging score functions and a novel regularization to improve approximation accuracy in high-dimensional stochastic systems.

Contribution

It proposes a new neural differential equation system for first- and second-order score functions, reformulating MFC problems as unconstrained optimizations with a regularization enforcing viscous HJB properties.

Findings

Effective in approximating solutions for MFC problems

Demonstrates accuracy on Wasserstein proximal operators and FP equations

Outperforms existing methods in key benchmarks

Abstract

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.