Variational conditional normalizing flows for computing second-order mean field control problems

TL;DR

This work introduces variational conditional normalizing flow (VCNF), a neural network-based variational algorithm for solving general MFC problems based on flow maps, and uses VCNF to model the Lagrangian formulation of the MFC problems.

Abstract

Mean field control (MFC) problems have vast applications in artificial intelligence, engineering, and economics, while solving MFC problems accurately and efficiently in high-dimensional spaces remains challenging. This work introduces variational conditional normalizing flow (VCNF), a neural network-based variational algorithm for solving general MFC problems based on flow maps. Formulating MFC problems as optimal control of Fokker-Planck (FP) equations with suitable constraints and cost functionals, we use VCNF to model the Lagrangian formulation of the MFC problems. In particular, VCNF builds upon conditional normalizing flows and neural spline flows, allowing efficient calculations of the inverse push-forward maps and score functions in MFC problems. We demonstrate the effectiveness of VCNF through extensive numerical examples, including optimal transport, regularized Wasserstein proximal operators, and flow matching problems for FP equations.

Figures (6)

• Figure 1: VCNF framework for solving generalized MFC problems.
• Figure 1: 2D Gaussian-to-Gaussian in OT problem. The source and target distributions are given by $\mathcal{N}\left(-33, 5110.5\right)$ and $\mathcal{N}\left(00, 1001\right)$ respectively. We visualize both the density evolution and several trajectories of the particles. In all the figures, the density heatmap is plotted for the domain $[-6, 3] \times [-3, 6]$. The four red curves represent the trajectories of four starting points $(-5, 3.5), (-5, 2.5), (-1, 3.5), (-1, 2.5)$ moving from upper left to lower right.
• Figure 2: 2D Gaussian mixture-to-Gaussian in OT problem. All trajectories start from the centers of the Gaussian mixture components, i.e. $(5, 0)$, $(3, 4)$, $(0, 5)$, $(-3, 4)$, $(-5, 0)$, $(-3, -4)$, $(0, -5)$, $(3, -4)$ and the plot domain is $[-7.5, 7.5] \times [-7.5, 7.5].$ The terminal Gaussian distribution is centered at the origin and all Gaussian distributions have covariance matrix $\mathbf{I}$.
• Figure 3: RWPO with initial condition of Gaussian distribution with covariance $\Sigma = 4\mathbf{I}$ and quadratic potential $V(\mathbf{x}) = \| \mathbf{x} \|_2^2/2$. We choose $T = 2, \beta = 1, \lambda = 200$ with t batch size $20$. The density on the domain $[-4, 4] \times [-4, 4]$ is visualized with trajectories starting from $(-3, -3), (-3, 3), (3, 3), (3, -3)$ that contract to the middle.
• Figure 4: RWPO with initial condition of Gaussian distribution with covariance $\Sigma = 4/5\mathbf{I}$ and double-well potential $V(\mathbf{x}) = ((x_1-1)^2+(x_2+1)^2)((x_1+1)^2 + (x_2-1)^2) / 4$. We choose $T= 2, \beta = 5, \lambda = 100$, and $N_t = 10$. The density on domain $[-2, 2] \times [-2, 2]$ is visualized and red points denote the trajectories starting from the boundary points $(-2, -2), (-2, 0), (-2, 2), (0, 2), (2, 2), (2, 0), (2, -2), (0, -2)$ and then contracting to two wells.

Theorems & Definitions (1)

• Proposition 1: Kernel solutions