Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Mo Zhou, Stanley Osher, Wuchen Li

TL;DR

This work advances mean field control by embedding first- and second-order score dynamics into a score-based neural ODE framework, enabling continuous-time evolution of log-densities and their derivatives along trajectories. By discretizing the neural ODEs into high-order normalization flows, the authors achieve efficient computation of the score and Hessian of the density, with a regularized MFC formulation that incorporates viscous HJB information. They develop both theoretical foundations and practical algorithms, including a regularized RWPO, FP-flow matching, and LQ/double-well benchmarks, demonstrating accurate density, velocity, and score recovery across dimensions. The approach offers a scalable, differentiable optimization pipeline for stochastic control problems in high dimensions, with potential impact on generative modeling, controlled diffusion, and inverse problems.

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.

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

TL;DR

This work advances mean field control by embedding first- and second-order score dynamics into a score-based neural ODE framework, enabling continuous-time evolution of log-densities and their derivatives along trajectories. By discretizing the neural ODEs into high-order normalization flows, the authors achieve efficient computation of the score and Hessian of the density, with a regularized MFC formulation that incorporates viscous HJB information. They develop both theoretical foundations and practical algorithms, including a regularized RWPO, FP-flow matching, and LQ/double-well benchmarks, demonstrating accurate density, velocity, and score recovery across dimensions. The approach offers a scalable, differentiable optimization pipeline for stochastic control problems in high dimensions, with potential impact on generative modeling, controlled diffusion, and inverse problems.

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.
Paper Structure (25 sections, 5 theorems, 134 equations, 9 figures, 5 tables, 4 algorithms)

This paper contains 25 sections, 5 theorems, 134 equations, 9 figures, 5 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Score-based neural ODE in continuous time
  3. Score-based normalization flows
  4. A generalization of normalizing flow
  5. Different algorithms for computing the score function
  6. Solving second-order MFC problems
  7. Formulation of MFC problem
  8. Modified HJB equation for MFC
  9. Numerical algorithms
  10. Numerical results
  11. Regularized Wasserstein proximal operator
  12. Flow matching for solving FP equations
  13. An LQ example with an entropy potential cost
  14. Double well potential
  15. Conclusion
  16. ...and 10 more sections

Key Result

Proposition 1

Functions $z_t$, $l_t$, $s_t$, and $H_t$ satisfy the following ODE dynamics. where $s_{it}$ and $f_i$ are the $i$-th component of $s_t$ and $f$ respectively.

Figures (9)

  • Figure 1: Numerical results for the RWPO problem. Left: cost functional through training with optimal cost in $1d$. Middle: density evolution through \ref{['eq:Lt_dynamic']} and comparison with true density in $1d$. Right: trajectories of $z_t$ and comparison with the stochastic dynamic in $2d$. The score dynamic demonstrates a structured behavior compared with the corresponding stochastic trajectory simulated by stochastic differential equations.
  • Figure 2: The first, second and third row: numerical results for RWPO in $1d$, $2d$, and $10d$. First column: cost functional through training. Second and third columns: first dimension of the velocity field at $t=0$ and $t=1$. Fourth column: the first row is the score function at $t=1$, the second and third rows are the density plots for the first dimension of the velocity field. Our numerical results accurately capture the true solutions.
  • Figure 3: Flow matching for OU process in $1$ dimension. Our algorithm accurately captures the density evolution of the state dynamic, including the change of mean and variance.
  • Figure 4: Flow matching for OU process in $2$ dimension. First line: particle evolution under the probability flow $z_t$ with trained velocity field. Second line: simulation of the stochastic OU process. The circles in the figure are centered at $\mu(t)$, with radius $\frac{1}{4} \sigma(t)$, $\frac{1}{2} \sigma(t)$, and $\sigma(t)$. Our algorithm accurately captures the density evolution of the state dynamic, including the change of mean and variance. The probability flow demonstrates more structured behavior compared with the stochastic OU process.
  • Figure 5: 2D flow matching double moon. First row: evolution of particles under the trained probability flow $z_t$ in the first stage $[0,0.2]$. Second row: evolution of particles under the overdamped Langevin dynamic in the first stage $[0,0.2]$. Third row: evolution of particles under the trained probability flow $z_t$ in the second stage $[0.2,0.4]$. Fourth row: evolution of particles under the overdamped Langevin dynamic in the second stage $[0.2,0.4]$. Our multi-stage splicing method captures the overdamped Langevin dynamic correctly.
  • ...and 4 more figures

Theorems & Definitions (12)

  • Proposition 1: Neural ODE system
  • Proposition 2: Information equality
  • Proposition 3
  • Corollary 1
  • proof
  • proof
  • proof
  • proof
  • proof
  • proof
  • ...and 2 more