# A Mean-Field Games Laboratory for Generative Modeling

Benjamin J. Zhang<EMAIL_ADDRESS>
Department of Mathematics and Statistics
University of Massachusetts Amherst
Amherst, MA 01003-9305, USA Markos A. Katsoulakis<EMAIL_ADDRESS>
Department of Mathematics and Statistics
University of Massachusetts Amherst
Amherst, MA 01003-9305, USA

###### Abstract

We demonstrate the versatility of mean-field games (MFGs) as a mathematical
framework for explaining, enhancing, and designing generative models. In
generative flows, a Lagrangian formulation is used where each particle
(generated sample) aims to minimize a loss function over its simulated path.
The loss, however, is dependent on the paths of other particles, which leads
to a competition among the population of particles. The asymptotic behavior of
this competition yields a mean-field game. We establish connections between
MFGs and major classes of generative flows and diffusions including
continuous-time normalizing flows, score-based generative models (SGM), and
Wasserstein gradient flows. Furthermore, we study the mathematical properties
of each generative model by studying their associated MFG’s optimality
condition, which is a set of coupled forward-backward nonlinear partial
differential equations. The mathematical structure described by the MFG
optimality conditions identifies the inductive biases of generative flows. We
investigate the well-posedness and structure of normalizing flows, unravel the
mathematical structure of SGMs, and derive a MFG formulation of Wasserstein
gradient flows. From an algorithmic perspective, the optimality conditions
yields Hamilton-Jacobi-Bellman (HJB) regularizers for enhanced training of
generative models. In particular, we propose and demonstrate an HJB-
regularized SGM with improved performance over standard SGMs. We present this
framework as an MFG laboratory which serves as a platform for revealing new
avenues of experimentation and invention of generative models.

Keywords: Generative modeling, mean-field games, Hamilton-Jacobi-Bellman
equation, normalizing flows, score-based generative models, Wasserstein
gradient flow, inductive bias

###### Contents

  1. 1 Introduction
    1. 1.1 A preview of results
  2. 2 Background on mean-field games
    1. 2.1 Optimality conditions
    2. 2.2 Potential formulation of MFGs
  3. 3 Continuous normalizing flows as solutions to MFGs
    1. 3.1 Continuous normalizing flows
    2. 3.2 Well-posedness and structure of canonical continuous normalizing flows
    3. 3.3 Variants of normalizing flows via MFGs
      1. 3.3.1 Optimal–transport normalizing flow
      2. 3.3.2 Boltzmann generators as MFGs
  4. 4 Score-based generative models as solutions to MFGs
    1. 4.1 Deriving SGM from an MFG
    2. 4.2 Score-matching objective through duality
    3. 4.3 MFGs unravel the curious mathematical serendipity of score-based generative modeling
    4. 4.4 Score-based probability flows as solutions of MFGs
    5. 4.5 Schrödinger bridges
      1. 4.5.1 MFGs, Schrödinger bridges, and a system of forward-backward heat equations
  5. 5 Generative models based on Wasserstein gradient flows as solutions to MFGs
    1. 5.1 Wasserstein gradient flows as solutions to MFG
    2. 5.2 Mean-field games with relaxation
    3. 5.3 Applications to generative modeling
  6. 6 Enhancing generative models with MFGs: HJB regularizers
    1. 6.1 Hamilton-Jacobi-Bellman regularizers
    2. 6.2 HJB regularizers for normalizing flows
      1. 6.2.1 Optimal-transport normalizing flows
      2. 6.2.2 Normalizing flows with bounded velocities
    3. 6.3 HJB regularizers for score-based generative modeling and beyond
    4. 6.4 A numerical demonstration: HJB-regularized SGM
    5. 6.5 Score-matching via PINNs
  7. 7 Inventing generative models using MFGs
    1. 7.1 Optimal transport BG: a corrected Boltzmann generator
    2. 7.2 Stochastic normalizing flows via MFGs
    3. 7.3 SGM and the cross-entropy regularization of optimal transport
    4. 7.4 Interpolating between MFGs: from Wasserstein gradient flows to normalizing flows and Wasserstein geodesics
  8. 8 A modular mean-field games framework for generative modeling
  9. 9 Discussion and outlook

##  1 Introduction

Generative models are powerful unsupervised learning tools that approximate a
target probability distribution via samples. These constructed models are
often used to generate new but approximate samples that statistically
reproduce and augment the original dataset. Generative models have recently
found extraordinary success in applications such as image generation, image
processing, audio generation, natural language processing, molecular modeling,
and more. Their success has led to a surge of research activity related to
formulating, parametrizing, and training new deep generative models. Three of
the most successful classes of generative models of the last decade are
generative adversarial nets (GANs) Goodfellow et al. (2020), denoising
diffusion probability models Ho et al. (2020), including score-based
generative models (SGM) Song et al. (2021b), and normalizing flows (NFs)
Rezende and Mohamed (2015); Chen et al. (2018); Grathwohl et al. (2018). These
three classes of models have spawned numerous competing variants in the
pursuit of inexpensive, high-quality generative models.

We focus on two groups of generative modeling techniques: continuous-time
generative models, and implicit generative models. Continuous normalizing
flows Chen et al. (2018) and score-based generative models Song et al. (2021b)
are examples of _continuous-time_ generative models, in which the model is
expressed as an ordinary or stochastic differential equation. The velocity
field is learned via data, such that the associated ODE or SDE transports
samples of a normal distribution to samples of the target distribution. In
contrast, _implicit_ generative models, based on Wasserstein gradient flows,
are less well-known particle system–based approaches for generative modeling
and sampling Arbel et al. (2019); Wang et al. (2021b); Gu et al. (2022). Here,
rather than learning an explicit map that generates samples from the target
distribution, new samples are instead generated through the Wasserstein
gradient flow of an empirical distribution toward the target distribution. The
Wasserstein gradient flow Santambrogio (2017) has also been used to study the
training of GANs Arbel et al. (2019); Mroueh et al. (2019). There is a
pervasive sense that flow and diffusion-based generative models have some
common fundamental structure and interconnections, and there have been several
attempts to connect these two classes of models together Albergo et al.
(2023); Zhang and Chen (2021). We will show that mean-field games provide a
natural framework for describing flow and diffusion-based generative models.

Mean-field game (MFG) theory is a mathematical field that studies the dynamics
and decision-making of individual agents in a large population Lasry and Lions
(2007); Guéant (2012). This emerging field has applications in many
disciplines including control theory, economics, finance, epidemiology, and
more. In this paper, we present a unifying mathematical framework based on
mean-field games for explaining, enhancing, and inventing generative models.
While the foundational formulations of the generative models we study in this
paper have different origins, we will demonstrate that these generative models
can be derived through a unifying mean-field games formalism.

Our primary observation in this paper is that mean-field games naturally
encode the characteristics of generative modeling with continuous-time models.
A mean-field game involves an objective function that models the action,
interaction, and terminal costs of an individual agent in relation to a
population of other agents. The solution to the MFG describes the action each
agent should take throughout space and time, while describing the evolution of
the population distribution over time. The optimality conditions of MFGs are a
pair of coupled PDEs: a Fokker-Planck (or continuity) equation, which
describes the evolution of the density, and a Hamilton-Jacobi-Bellman (HJB)
equation, which describes the optimal velocity field.

Meanwhile, continuous-time generative models have a common set of ingredients.
Given samples from some target distribution, an optimal velocity field that
evolves a simple reference distribution to the target distribution in finite
time is learned. An optimization problem over the space of probability
distributions is solved to learn the optimal velocity field. The evolution of
the density is determined by a transport PDE (the continuity equation or
Fokker-Planck equation), which is intractable to solve beyond low dimensional
settings. Instead, the Lagrangian formulation is considered so that the
generative modeling task becomes an optimization problem where each individual
particle (generated sample) evolves according to its location while minimizing
some objective function over the path it takes. This optimization problem
appears to be an optimal control problem; it is, however, a mean-field game
since the objective function depends on other generated samples (particles).
In the language of mean-field games, the Fokker-Planck describes the family of
generative models, while the solution of the HJB characterizes the optimal
generative model.

We will derive the major classes of generative models with the MFG formalism
and culminate in a sort of laboratory (see Table 3) in which researchers can
experiment with existing and new generative modeling approaches. We also
create several new examples of generative models as demonstration of the
potential uses of the MFG laboratory.

##### Our contributions:

We present a unifying mean-field games framework for explaining, enhancing,
and inventing generative models. By _explaining_ how generative models can be
derived from mean-field games, we bring attention to the richness of the MFG
for designing and analyzing generative models.

  * •

Explaining generative models: Formulating generative flows as MFGs reveals a
_common mathematical structure_ through their optimality conditions, i.e., a
backwards HJB and a forward Fokker-Planck equation (FPE). This
backward–forward structure was already observed as a key feature in denoising
diffusion models. Our MFG formulation shows that it is shared by _all_ flow
and diffusion–based generative models, and delineates the _generative_ and
_discriminative_ features required to describe the flows. That is, the FPE
describes the generative model while the HJB identifies the optimal velocity
field. The mathematical structure described by the mean-field game identifies
the inductive bias of generative flows. In particular, we formulate
normalizing flow (Theorem 2), score-based generative models (Theorem 7), and
Wasserstein gradient flows (Theorem 14) in terms of MFGs. Another key result
is that the MFG optimality conditions characterize the well-posedness of
optimization problems used to train generative models. For example, in Section
3.2, we study the ill-posedness of the canonical normalizing flows framework
using the theory of HJB PDEs, and ways to make the method well-posed.

  * •

Enhancing generative models: MFGs inform structure of generative flows that
may be used to accelerate their training. The optimality conditions explicitly
relate the optimal velocities to the gradient of the solution of the HJB.
Therefore, even without solving the HJB explicitly, we know that we can
_restrict_ the generative vector fields to just gradients of a scalar
function. Moreover, in Section 6 we construct regularization functionals that
can be derived from the optimality conditions which may further accelerate
learning of the optimal generative model.

  * •

Inventing generative models: This unifying approach establishes a platform for
systematically building new generative models. In Section 7 we give examples
of new generative modeling formulations using MFGs.

Our work culminates to Table 3 in Section 8, which not only classifies major
generative models, but also provides a laboratory for the experimentation,
investigation, and invention of new generative models. We emphasize, however,
that not all models developed using the laboratory will necessarily have good
practical performance. This framework allows one to formulate generative
models with well-defined optimization problems and inform their associated
structure, however one must still proceed with numerical and algorithmic
investigation to make their implementation practical.

##### Related work:

The connections between mean-field games and machine learning have been
previously investigated. The use of generative modeling tools, such as
normalizing flows, to solve mean-field games has been well-explored Ruthotto
et al. (2020); Huang et al. (2022). Moreover, connections between
reinforcement learning and mean-field games have been explored as well Guo et
al. (2022). Studying score-based generative models with stochastic optimal
control tools is discussed in Berner et al. (2022). The use of partial
differential equations to enhance the training of score-based generative
models was applied in Lai et al. (2023). Recently, Albergo and Vanden-Eijnden
(2022); Albergo et al. (2023) have also unified normalizing flows with
diffusion models and the Schrödinger bridge through _stochastic interpolants_.
To our knowledge, however, there has not been a comprehensive unifying
approach of using the MFG formalism to describe several classes of generative
models.

###  1.1 A preview of results

We preview the main results of this paper. In Section 2, we provide a more
thorough exposition of mean-field games, but we provide a simplified
presentation here to discuss the main ideas of our paper.

Let $X$ be a $\mathbb{R}^{d}$–valued random variable with unknown data
distribution $\pi$ and density $\pi(x)$. Denote $\rho_{0}$ be a standard
normal distribution with density $\rho_{0}(x)$. We suppose we have only
samples $\\{x_{i}\\}_{i=1}^{N}$ from $\pi$. The generative modeling problem is
the task of producing more samples from $\pi(x)$.

The potential formulation of mean-field games (MFG) is an infinite dimensional
optimization problem over velocity fields $v\in\mathcal{C}^{2}$,
$v:\mathbb{R}^{d}\times\mathbb{R}\to\mathbb{R}^{d}$, and densities
$\rho(\cdot,t)\in\mathcal{P}(\mathbb{R}^{d})$ with finite second moments for
some finite time interval $t\in[0,T]$. Broadly speaking, the objective
function is decomposed into a sum of three types of cost functions, $L$ to
model action costs, $\mathcal{I}$ to model interaction costs, and
$\mathcal{M}$ to model terminal costs. The continuity equation constraint
models the dynamics of the density of the agents under a particular velocity
field. There is a correspondence between the continuity equation and the
particle dynamics, which depend on whether the dynamics are deterministic or
stochastic. The MFG is defined as

$\displaystyle\inf_{v,\rho}\left\\{\mathcal{M}(\rho(\cdot,T))+\int_{0}^{T}\mathcal{I}(\rho(\cdot,t))\,\text{d}t+\int_{0}^{T}\int_{\mathbb{R}^{d}}L(x,v(x,t))\rho(x,t)\,\text{d}x\,\text{d}t\right\\}$
(1) s.t. $\displaystyle\,\,\frac{\partial\rho}{\partial
t}+\nabla\cdot(v\rho)=\frac{\sigma^{2}}{2}\Delta\rho,\,\rho(x,0)=\rho_{0}(x)\,.$

The optimality conditions for the MFG are given by a coupled system of a
continuity (Fokker-Planck) equation for $\rho$ which corresponds to the
optimal generative model, and a (backward) Hamilton-Jacobi-Bellman (HJB)
equation which yields the optimal velocity field $v^{*}$.

We will derive canonical formulations of generative models using the potential
formulation of mean-field games. Depending on how cost functions
$L,\mathcal{I}$, and $\mathcal{M}$, are chosen, the solutions of the
corresponding MFGs will be different classes of generative models. Aspects of
the derivations and connections have been noted in the literature for
particular cases, such as the optimal transport regularized normalizing flow
Onken et al. (2021); Ruthotto et al. (2020); Huang et al. (2022), and the
stochastic optimal control perspective of score-based generative modeling
Berner et al. (2022). Here, we provide a comprehensive unifying approach,
using the full MFG formalism to describe several classes of fundamental
generative models. In Tables 1 and 2, we overview the main contribution of
this paper for three well-studied classes of generative models. Eventually, we
will build to Table 3, by demonstrating how numerous generative modeling
variants can be described via the MFG framework. In Figure 1, we show a
schematic illustrating how the mean-field games perspective relates and
provides insight to the current understanding of generative models.

Generative model | $\mathcal{M}(\rho)$ | $\mathcal{I}(\rho)$ | $L(x,v)$
---|---|---|---
Normalizing flow | $\mathbb{E}_{\rho}\left[\log\rho-\log\rho_{ref}\right]$ | 0 | 0
Score-based | $-\mathbb{E}_{\rho}\log\pi$ | 0 | $\frac{1}{2}|v|^{2}-\nabla\cdot f$
Wasserstein Gradient Flows ($\epsilon\to 0$) | $\mathcal{F}(\rho)e^{-T/\epsilon}$ | $\frac{e^{-t/\epsilon}}{\epsilon}\mathcal{F}(\rho)$ | $\frac{1}{2}|v|^{2}e^{-t/\epsilon}$
Table 1: Generative models and corresponding mean-field game cost functions. See a more complete version in Table 3. Generative model | Dynamics | $H(x,p)$ | Optimal $v^{*}$
---|---|---|---
Normalizing flow | $\,\text{d}x=v\,\text{d}t$ | $\sup_{v\in K\subset\mathbb{R}^{d}}-p^{\top}v$ | $-\nabla_{p}H(x,\nabla U)$
Score-based | $\,\text{d}x=\left[f+\sigma v\right]\,\text{d}t+\sigma\,\text{d}W_{t}$ | $-f^{\top}p+\frac{|\sigma p|^{2}}{2}+\nabla\cdot f$ | $\sigma\nabla\log\eta_{t}$
Wasserstein GF ($\epsilon\to 0$) | $\,\text{d}x=v\,\text{d}t$ | $\frac{1}{2}e^{t/\epsilon}|p|^{2}$ | $-\frac{\delta\mathcal{F}}{\delta\rho}$
Table 2: Generative models, their corresponding dynamics, Hamiltonians, and
optimal velocity fields. Figure 1: Flow chart describing how mean-field games
are related to and provide new insights for flow-based generative modeling.
Blue boxes denote the current understanding of flow-based generative models.
Green boxes are the new perspective provided by mean-field games.

##  2 Background on mean-field games

We review formulations of mean-field games and their optimality conditions.
This section is based on Lasry and Lions (2007); Bensoussan et al. (2013).
Mean-field games model the dynamics of individual rational agents whose
actions depend on the dynamics of a large population of other agents. When the
trajectories and actions take values in $\mathbb{R}^{d}$, the MFG is a
differential game, whose solution determines the optimal action of each
individual agent. The optimality conditions are a system of partial
differential equations, the Fokker-Planck (continuity) and Hamilton-Jacobi-
Bellman (HJB) equations, which characterize the solution of the MFG. The
solution of the optimality conditions is often referred to as a _Nash
equilibrium_.

Denote $\mathcal{P}(\mathbb{R}^{d})$ to be the space of probability
distributions on $\mathbb{R}^{d}$ with finite second moment, and suppose the
evolution of an agent $x(s)\in\mathbb{R}^{d}$ over the time interval
$s\in[0,T]$ is given by the stochastic differential equation

$\displaystyle\begin{dcases}\,\text{d}x(s)&=v(x(s),s)\,\text{d}s+\sigma(x(s),s)\,\text{d}W(s)\\\
x(0)&=x_{0}\end{dcases}$ (2)

where $v:\mathbb{R}^{d}\times\mathbb{R}\to\mathbb{R}^{d}$ is a vector field
that models the action of the agent at time $s$ and position $x(s)$,
$\sigma(x(s),s)$ is the diffusion coefficient, and $W(s)$ is a standard
Brownian motion on $\mathbb{R}^{d}$. The velocity field (action) at $(x,t)$ is
determined by minimizing the objective function
$J:\mathcal{C}^{2}(\mathbb{R}^{d})\times\mathcal{P}(\mathbb{R}^{d})\to\mathbb{R}$,

$\displaystyle
J_{x,t}(v,\rho)=\mathbb{E}\left[\int_{t}^{T}L(x(s),v(x(s),s))+I(x(s),\rho(x(s),s))\,\text{d}s+M(x(T),\rho(x(T),T))\right],$
(3)

where $L:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$ is the running cost
incurred by the agent by performing action $v(x(s),s)$,
$I:\mathbb{R}^{d}\times\mathcal{P}(\mathbb{R}^{d})\to\mathbb{R}$ is the
running cost due to interaction with the density of all the agents
$\rho(\cdot,s)\in\mathcal{P}(\mathbb{R}^{d})$, and
$M:\mathbb{R}^{d}\times\mathcal{P}(\mathbb{R}^{d})\to\mathbb{R}$ is the
terminal cost which depends on the final position of the agent and the
population density. The expectation is taken over the path-space measure of
(2) with initial condition $(x,t)$ and over time interval $s\in[t,T]$. Note
that the running costs may explicitly depend on time, but for simplicity of
notation, we omit the time argument. We consider a time-dependent example in
Section 5. Let $\mathbf{D}(x,t)=\frac{1}{2}\sigma(x,t)\sigma(x,t)^{\top}$.
Given a velocity field $v=v(x,t)$, the evolution of the population density is
given by the continuity (Fokker-Planck) equation

$\displaystyle\frac{\partial\rho}{\partial
s}+\nabla\cdot(v\rho)=\nabla^{2}:\mathbf{D}\rho$ (4)
$\displaystyle\rho(x,0)=\rho_{0}(x)$

Here, $\mathbf{A}:\mathbf{B}=\sum_{ij}\mathbf{A}_{ij}\mathbf{B}_{ij}$, and
$(\nabla^{2})_{ij}=\frac{\partial^{2}}{\partial x_{i}\partial x_{j}}$. This
formulation also holds when the agent dynamics in (2) are deterministic, i.e.,
when $\sigma\coloneqq 0$.

###  2.1 Optimality conditions

The solution of a mean-field game is a pair
$({v^{*}}(\cdot,t),{\rho^{*}}(\cdot,t))$, referred to as a Nash equilibrium,
such that for all $v$,

$\displaystyle J_{x,t}({v^{*}},{\rho^{*}})\leq J_{x,t}(v,{\rho^{*}}).$ (5)

The optimality conditions characterizing this pair is found by considering the
value function

$\displaystyle U(x,t)=\inf_{v,\rho}J_{x,t}(v,\rho),$ (6)

which solves a Hamilton-Jacobi-Bellman (HJB) equation that is coupled to the
continuity equation. The optimality conditions are

$\displaystyle\begin{dcases}-\frac{\partial U}{\partial t}+H(x,\nabla
U)-\mathbf{D}:\nabla^{2}U=I(x,\rho(x,t))\\\ \frac{\partial\rho}{\partial
t}-\nabla\cdot(\rho\nabla_{p}H(x,\nabla U))=\nabla^{2}:\mathbf{D}\rho\\\
U(x,T)=M(x,\rho({x},T))\\\ \rho(x,0)=\rho_{0}(x).\end{dcases}$ (7)

Here, the Hamiltonian is defined as

$\displaystyle H(x,p)=\sup_{v}\left[-p^{\top}v-L(x,v)\right],$ (8)

and the optimal control, ${v^{*}}(x,t)$, is given by the maximizer, and in
this case is equal to

$\displaystyle{v^{*}}(x,t)=-\nabla_{p}H(x,\nabla U(t,x)).$ (9)

Since this relation holds at optimality, it is common to replace $v$ with
${v^{*}}$ in (4) to explicitly couple the continuity and HJB equations. Note
that the HJB equation is solved _backwards_ in time while the continuity
equation is solved _forwards_ time as the latter has a terminal condition, and
the former has an initial condition. Moreover, it is important to note that
this is not a standard optimal control problem as the HJB equation cannot be
solved by itself since it is coupled with the continuity equation. For details
about the derivation of the optimality conditions, see Lasry and Lions (2007)
and the references therein.

###  2.2 Potential formulation of MFGs

Solving the coupled system of PDEs in (7) via traditional numerical methods
for PDEs is intractable beyond low dimensional settings. Meanwhile, solving
the MFG directly as presented in (2) and (3) is still quite computationally
taxing as it requires solving an infinite dimensional optimization problem for
each point in time and space. In the celebrated paper Lasry and Lions (2007),
an alternative _potential formulation_ is devised for a certain class of MFGs
so that the optimal controls and densities can be determined simultaneously
for all $(x,t)\in\mathbb{R}^{d}\times\mathbb{R}$.

Suppose there exist functionals
$\mathcal{M}:\mathcal{P}(\mathbb{R}^{d})\to\mathbb{R}$ and
$\mathcal{I}:\mathcal{P}(\mathbb{R}^{d})\to\mathbb{R}$ such that

$\displaystyle
I(x,\rho)=\frac{\delta\mathcal{I}(\rho)}{\delta\rho}(x),\,\,\,\,M(x,\rho)=\frac{\delta\mathcal{M}(\rho)}{\delta\rho}(x).$
(10)

For a test function $\phi\in L^{2}(\mathbb{R}^{d})$ and some density function
$\rho\in\mathcal{P}(\mathbb{R}^{d})$, recall that the variational derivative
of functional $\mathcal{F}$ is given by

$\displaystyle\lim_{h\to
0}\frac{\mathcal{F}(\rho+h\phi)-\mathcal{F}(\rho)}{h}=\int_{\mathbb{R}^{d}}\frac{\delta\mathcal{F}}{\delta\rho}\phi(x)\,\text{d}x.$

The potential formulation of the MFG (3) is given by

$\displaystyle\inf_{v,\rho}\mathcal{J}(v,\rho)=$
$\displaystyle\inf_{v,\rho}\left\\{\mathcal{M}(\rho(\cdot,T))+\int_{0}^{T}\mathcal{I}(\rho(\cdot,t))\,\text{d}t+\int_{0}^{T}\int_{\mathbb{R}^{d}}L(x,v(x,t))\rho(x,t)\,\text{d}x\,\text{d}t\right\\}$
(11) s.t. $\displaystyle\,\,\frac{\partial\rho}{\partial
t}+\nabla\cdot(v\rho)=\nabla^{2}:\mathbf{D}\rho,\,\rho(x,0)=\rho_{0}(x).$

The corresponding optimality conditions (7) are,

$\displaystyle\begin{dcases}-\frac{\partial U}{\partial t}+H(x,\nabla
U)-\mathbf{D}:\nabla^{2}U=I(x,\rho)=\frac{\delta\mathcal{I}(\rho)}{\delta\rho}(x)\\\
\frac{\partial\rho}{\partial t}-\nabla\cdot(\nabla_{p}H(x,\nabla
U)\rho)=\nabla^{2}:\mathbf{D}\rho\\\
U(x,T)=M(x,\rho)=\frac{\delta\mathcal{M}(\rho(\cdot,T))}{\delta\rho}(x)\\\
\rho(x,0)=\rho_{0}(x).\end{dcases}$ (12)

The optimization problem given in this formulation is often _implicitly_
considered in many machine learning tasks, such as density estimation and
generative modeling. We will demonstrate the connections between the potential
MFG and generative models clearly. In the following sections, we will use the
potential MFG formulation (11) to establish for a unified perspective of
normalizing flows, score-based generative modeling, Wasserstein gradient
flows, and their variants. It will also provide a framework for extending and
enhancing these classes of generative models.

##  3 Continuous normalizing flows as solutions to MFGs

A normalizing flow (NF) is an _invertible_ transport map
$f:\mathbb{R}^{d}\to\mathbb{R}^{d}$ that a transforms target distribution
$\pi$ to a standard normal $\rho_{ref}$ Rezende and Mohamed (2015); Tabak and
Vanden-Eijnden (2010); Tabak and Turner (2013). By learning such a map $f$
through samples of $\pi$, one can then produce more realizations from the
target distribution by inverting the map on sample points of the normal
distribution. We say reference distribution $\rho_{ref}$ is the _pushforward_
distribution of $\pi$ under map $f$, and their densities $\pi(x)$, and
$\rho_{ref}(x)$ are related by a change of variables formula

$\displaystyle\log\pi(x)=\log\rho_{ref}(f(x))+\log\det\mathbf{J}_{f}(x)$ (13)

where $\mathbf{J}_{f}(x)$ denotes the Jacobian of $f$. For generative
modeling, the map is typically learned by maximizing the log-likelihood of the
pullback distribution defined by the approximate map $f_{\theta}$ with respect
to the data distribution over a parametrized space of functions, i.e., by
solving the variational problem

$\displaystyle\max_{\theta}\mathbb{E}_{\pi}\left[\log\pi_{\theta}(x)\right]=\max_{\theta}\mathbb{E}_{\pi}\left[\log\rho_{ref}(f_{\theta}(x))+\log\det\mathbf{J}_{f_{\theta}}(x)\right]$
(14)

for parameters $\theta\in\Theta$.

The map $f_{\theta}$ is often parameterized as a composition of simple
invertible functions, which makes the map amenable to approximation via neural
networks. Each layer of the neural network represents one of the composed
functions. In practice, the simple invertible functions are chosen depending
on the particular application and so that the determinant Jacobian can be
computed quickly. There is a rich body of literature outlining different
function choices. See Papamakarios et al. (2021); Kobyzev et al. (2020) and
the references therein for a review of parametrization choices.

###  3.1 Continuous normalizing flows

A _continuous normalizing flow_ (CNF) is a particular type of normalizing flow
that can be conceptualized as a standard NF as the number of layers in the
neural network tends to infinity Chen et al. (2018). The transport map is
described as the solution of a differential equation, and the velocity field
of the differential equation is learned from data rather than the entire flow
map. The goal then is to learn a $v_{\theta}(x,t)$ such that we have an ODE
model $\frac{\,\text{d}x}{\,\text{d}t}=v_{\theta}(x(t),t)$, $x(0)\sim\pi$, and
$x(T)\sim\rho(\cdot,T)=\rho_{ref}$, where $\rho(x,t)$ satisfies the continuity
equation (4). Running the normalizing flow reverse in time then produces a
generative model that flows samples from $\rho_{ref}$ to new samples from
$\pi$.

A computable change of variable formula also exists for continuous normalizing
flows. Assuming that $v$ is uniformly Lipschitz in $x$ and continuous in $t$.
The change in log probability along a trajectory $x(t)$ is given by the
instantaneous change of variables formula Chen et al. (2018)

$\displaystyle\frac{D\log\rho(x(t),t)}{Dt}=-\nabla\cdot
v(x,t)\coloneqq-\sum_{i=1}^{d}\frac{\partial v_{i}}{\partial x_{i}}.$ (15)

The likelihood can be estimated by simply integrating (15) along with
$\,\text{d}x/\,\text{d}t=v(x(t),t)$ simultaneously. A proof of this formula is
given in Chen et al. (2018); however it is a direct consequence of the
definition of _material derivative_ and the continuity equation, as we discuss
in Remark 1.

######  Remark 1 (Instantaneous change of variables and the material
derivative).

The instantaneous change of variables formula Chen et al. (2018) is a direct
consequence of a standard result from fluid mechanics. In the Lagrangian
description of fluid mechanics, the material derivative describes the rate of
change of a quantity along a single flow trajectory Munson et al. (2013). For
a scalar quantity $\phi(x,t)$ and flow velocity $v(x,t)$, the material
derivative is defined as

$\displaystyle\frac{D\phi}{Dt}\coloneqq\frac{\partial\phi}{\partial
t}+v\cdot\nabla\phi.$ (16)

Therefore, the change of variables formula is simply the material derivative
applied to the log-density of the normalizing flow. Observe that

$\displaystyle\frac{D}{Dt}\log\rho(x(t),t)$
$\displaystyle=\frac{1}{\rho(x,t)}\frac{\partial\rho(x,t)}{\partial
t}+\frac{v(x,t)}{\rho(x,t)}\cdot\nabla\rho(x,t)$
$\displaystyle=-\frac{1}{\rho}\nabla\cdot(v\rho)+\frac{v\cdot\nabla\rho}{\rho}$
$\displaystyle=-\nabla\cdot v.$

Normalizing flows easily fit into the mean-field games framework. Our first
result is a general normalizing flows formulation in terms of mean-field
games, and is given in Theorem 2.  Given a reference distribution
$\rho_{ref}$, some running cost $L(x,v)$, $\mathcal{I}(\rho)\coloneqq 0$, and
$\mathcal{M}(\rho)=\mathcal{D}_{KL}(\rho\|\rho_{ref})$, with the initial
density being the target data distribution $\pi(x)$, many variants of
normalizing flows can be derived. Here
$\mathcal{D}_{KL}(\rho\|\rho^{\prime})=\mathbb{E}_{\rho^{\prime}}\left[\log\frac{\rho}{\rho^{\prime}}\right]$
is the Kullback-Leibler divergence.

######  Theorem 2.

Let $\pi$ be the target distribution. Generative models based on normalizing
flows are solutions to the mean-field game

$\displaystyle\min_{v,\rho}\left\\{{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\mathcal{D}_{KL}(\rho(\cdot,T)\|\rho_{ref})}+\int_{0}^{T}\int_{\mathbb{R}^{d}}L(x,v)\rho(x,t)\,\text{d}x\,\text{d}t:\frac{\partial\rho}{\partial
t}+\nabla\cdot\left(\rho
v\right)=0,\,\rho(x,0)={\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\pi(x)}\right\\},$
(17)

with dynamics

$\displaystyle\frac{\,\text{d}x}{\,\text{d}t}=v(x(s),s).$ (18)

Moreover, for sets $K\subseteq\mathbb{R}^{d}$, the solution of the mean-field
game satisfies the optimality conditions for $H(x,p)=\sup_{v\in
K}\left[-p^{\top}v-L(x,v)\right],$

$\displaystyle\begin{dcases}-\frac{\partial U}{\partial t}+H(x,\nabla U)=0\\\
\frac{\partial\rho}{\partial t}-\nabla\cdot(\rho\nabla_{p}H(x,\nabla U))=0\\\
U(x,T)={\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}1+\log\frac{\rho(x,T)}{\rho_{ref}(x)}},\,\rho(x,0)={\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\pi(x)},\end{dcases}$
(19)

where the optimal velocity field of the continuous normalizing flow has the
representation formula

$\displaystyle{v^{*}}(x,t)=\arg\sup_{v}\left\\{-p^{\top}v-L(x,v)\right\\}=-\nabla_{p}H(x,\nabla
U).$ (20)

Choosing the running costs to be identically zero, $L=0$, $\mathcal{I}=0$, and
the terminal cost to be the KL divergence from the density $\rho(\cdot,T)$ to
the reference distribution,
$\mathcal{M}(\rho)=\mathcal{D}_{\text{KL}}(\rho\|\rho_{ref})$, we will recover
the standard normalizing flows objective in Chen et al. (2018); Grathwohl et
al. (2018). This is the standard KL divergence minimization problem, and its
relationship to maximum likelihood estimation is very well-noted and
understood.

Proof  We perform the following computations, while always implicitly assuming
sufficient smoothness of all functions. First note that by a change of
variables, the terminal conditions can be re-expressed as follows

$\displaystyle\mathcal{D}_{KL}(\rho(x,T)\|\rho_{ref}(x))=\mathcal{D}_{KL}(\pi(x)\|\tilde{\rho}(x,0)),$
(21)

where $\tilde{\rho}(x,0)$ is the density of the inverse flow at time $t=0$
given that $\tilde{\rho}(x,T)=\rho_{ref}(x)$. Next observe that

$\displaystyle\mathcal{D}_{KL}(\pi(x)\|\tilde{\rho}(x,0))=\mathbb{E}_{\pi}\left[\log\pi(x)\right]-\mathbb{E}_{\pi}\left[\log\tilde{\rho}(x,0)\right].$

Furthermore, $\log\tilde{\rho}(x,0)$ can be expressed in terms of the
reference density $\rho_{ref}(x)$ through the continuity equation and the
definition of the material derivative. Integrating along a trajectory $x(t)$,
we have that for $x(T)\sim\pi$,

$\displaystyle\log\tilde{\rho}(x(0),0)$
$\displaystyle=\log\rho_{ref}(x(T))-\int_{T}^{0}\nabla\cdot
v(x(s),s)\,\text{d}s$ (22) $\displaystyle x(0)$
$\displaystyle=x(T)+\int_{T}^{0}v(x(s),s)\,\text{d}s.$

Note that for $L(x,v)=0$, the resulting objective function is precisely the
normalizing flows objective as described in Chen et al. (2018); Grathwohl et
al. (2018). Observe that (17) is  now equivalent to

$\displaystyle\min_{v}\left\\{-\mathbb{E}_{\pi}\left[\log\rho_{ref}(x)-\int_{T}^{0}\nabla\cdot
v(x(s),s)\,\text{d}s\right]:x(s)=x+\int_{T}^{s}v(x(s^{\prime}),s^{\prime})\,\text{d}s^{\prime},\,x\sim\pi\right\\}.$
(23)

As for the optimality conditions, observe that for test function $\chi(x)$
where $\int\chi(x)\,\text{d}x=0$,

$\displaystyle\int\frac{\delta\mathcal{M}}{\delta\rho}\chi(x)\,\text{d}x$
$\displaystyle=\lim_{\epsilon\to
0}\frac{1}{\epsilon}\int\left(\log\frac{\rho(x)+\epsilon\chi(x)}{\rho_{ref}(x)}\right)(\rho(x)+\epsilon\chi(x))-\left(\log\frac{\rho(x)}{\rho_{ref}(x)}\right)\rho(x)\,\text{d}x$
$\displaystyle=\lim_{\epsilon\to
0}\frac{1}{\epsilon}\int\left(\log\left[\frac{\rho(x)+\epsilon\chi(x)}{\rho(x)}\right]\rho(x)+\epsilon\log\left[\frac{\rho(x)+\epsilon\chi(x)}{\rho_{ref}(x)}\right]\chi(x)\right)\,\text{d}x$
$\displaystyle=\lim_{\epsilon\to
0}\int\left(\rho(x)\frac{\log(\rho(x)+\epsilon\chi(x))-\log\rho(x)}{\epsilon}+\log\left[\frac{\rho(x)+\epsilon\chi(x)}{\rho_{ref}(x)}\right]\chi(x)\right)\,\text{d}x$
$\displaystyle=\int\left(1+\log\frac{\rho(x)}{\rho_{ref}(x)}\right)\chi(x)\,\text{d}x$

so $M(x,\rho)=1+\log\frac{\rho(x)}{\rho_{ref}(x)}$. Lastly the Hamiltonian and
optimality conditions are clear from the potential formulation of MFG in (12),
see also Lasry and Lions (2007).

######  Remark 3.

Even though there is no running cost function due to interaction,
$\mathcal{I}(x,\rho)$, (17) is still not an optimal control problem since the
continuity and HJB equations are coupled by the terminal condition. This
implies that normalizing flows are necessarily solutions of mean-field games.

###  3.2 Well-posedness and structure of canonical continuous normalizing
flows

In this section we focus on the original formulation of continuous normalizing
flows as discussed in Chen et al. (2018); Grathwohl et al. (2018). While, in
practice, training a normalizing flow does not require the MFG perspective,
there is, however, significant insight that can be gained by studying this
objective function in the context of MFGs. In particular, we can study the
optimality conditions of the corresponding HJB equation to understand the
_well-posedness_ of the optimization problem and the _structure_ of the
optimal velocity field $v^{*}(x,t)$.

######  Corollary 4.

Let $L=0$, and $K\subset\mathbb{R}^{d}$ be a compact convex set. Then the
optimizer of $\sup_{v\in K}-p^{\top}v$ is attained on the boundary of $K$.
Moreover, when $K$ is a ball of radius $c$, then the optimal velocity field is

$\displaystyle v^{*}(x,t)=-c\frac{\nabla U(x,t)}{|\nabla U(x,t)|},$ (24)

where $U(x,t)$ solves the level-set equation Osher and Fedkiw (2001); Sethian
(1999) in the system of PDEs

$\displaystyle\begin{dcases}\frac{\partial U}{\partial t}-c|\nabla U|=0\\\
\frac{\partial\rho}{\partial t}-\nabla\cdot\left(c\rho\frac{\nabla U}{|\nabla
U|}\right)=0\\\
U(x,T)={\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}1+\log\frac{\rho(x,T)}{\rho_{ref}(x,T)},\,\rho(x,0)=\pi(x).}\end{dcases}$
(25)

Proof  Following the notation of Theorem 2, the Hamiltonian of the system is
$H(x,p)=\sup_{v\in K}-p^{\top}v.$ Since $p^{\top}v$ is a convex function in
$v$ and $K$ is convex and compact, by basic convex analysis, the optimizer
must lie on the boundary Boyd et al. (2004). Therefore, if $K$ is ball of
radius $c$, then the optimizer is $v=-cp/|p|$ so $H(x,p)=c|p|$, and by (7), we
obtain the desired optimality conditions with

$\displaystyle v^{*}(x,t)=-c\frac{\nabla U(x,t)}{|\nabla U(x,t)|}.$ (26)

Corollary 4 shows that for the standard continuous normalizing flow (CNF)
formulation Chen et al. (2018); Grathwohl et al. (2018), the Hamiltonian
cannot be defined properly without additional assumptions on the set of
feasible velocity fields $K\subset\mathbb{R}^{d}$. Specifically, $K$ is
necessarily a bounded set for the Hamiltonian to be finite and for the
normalizing flows optimization problem to be well-posed. Furthermore, if $K$
is a convex set, the optimal control must lie on the boundary $\partial K$. We
can interpret this as a form of _bang-bang_ control, in which the optimal
control always lies on the boundary of the feasible set Evans (2005).

######  Remark 5.

The typical assumption when learning CNFs is that the velocity fields only
need to be Lipschitz continuous so that the resulting ODE has a solution. We
argue that our result here shows that the velocity field must come from a
bounded set of feasible controls and that successful implementations of CNFs
implicitly enforce this constraint. In fact, Verine et al. (2021) states that
in practice, most normalizing flows are bi-Lipschitz, meaning that both the
map and its inverse are Lipschitz. Corollary 4 describes the precise
mathematical justification for this observation.

######  Remark 6 (Well-posedness of normalizing flows through MFG).

Learning neural ODEs through empirical risk minimization is an ill-posed
problem Ott et al. (2020). We argue that the MFG formalism may be one approach
to study the well-posedness of normalizing flow. Indeed, taking advantage of
the MFG optimality conditions, the normalizing flow variational problem is
well-posed only if its corresponding MFG PDE system is well-posed Lasry and
Lions (2007); Bensoussan et al. (2013). Moreover, we may even study the well-
posedness of certain GANs formulations based on normalizing flows, such as the
Flow-GAN Grover et al. (2018). If the generator of a GAN is a continuous
normalizing flow, then we may study them as an MFG with different choices of
terminal cost.

###  3.3 Variants of normalizing flows via MFGs

With the MFG framework for normalizing flows at hand in Theorem 2, we are able
to derive known variants of normalizing flows, re-write them as MFG, and
analyze them based on their optimality conditions and corresponding PDEs.

####  3.3.1 Optimal–transport normalizing flow

A special MFG that is well-studied is the case when the running cost is a term
that can be interpreted as kinetic energy $L(x,v)=\frac{1}{2}|v|^{2}$. This
yields a Hamiltonian that is quadratic in the co-state as
$H(x,p)=\frac{1}{2}|p|^{2}$, which together with Theorem 2 implies that the
optimal control is simply the gradient of the value function:
$v^{*}(t,x)=-\nabla U(t,x)$. Not only does this special case simplify and
provide additional structure to the problem, it also reveals connections to
optimal transport and other classes of generative models. In particular,
quadratic running cost functions can be used to regularize the training of
normalizing flows Onken et al. (2021); Huang et al. (2022). The addition of a
nonzero running cost to the normalizing flow MFG is also another approach to
make the normalizing flow optimization problem well-posed.

In Onken et al. (2021), the authors introduce the optimal transport flow (OT-
flow), which trains normalizing flows with the additional quadratic running
cost. They presented their method as combining optimal transport with
normalizing flows. Through numerical examples, they found improved regularity
of the normalizing flow models in that the resulting trajectories have
straighter solutions, which are easier to integrate over time. The OT-flow is
the following potential MFG:

$\displaystyle\inf_{v,\rho}\left\\{{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\mathcal{D}_{KL}(\rho(\cdot,T)\|\rho_{ref})}+\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)\,\text{d}x\,\text{d}t:\frac{\partial\rho}{\partial
t}+\nabla\cdot(v\rho)=0,\,{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\rho(x,0)=\pi(x)}\right\\}.$
(27)

The Hamiltonian is $H(x,p)=\frac{1}{2}|p|^{2}$, so by Theorem 2,
$v(x,t)=-\nabla U(x,t)$ with optimality conditions

$\displaystyle\begin{dcases}&-\frac{\partial U}{\partial t}+\frac{1}{2}|\nabla
U|^{2}=0\\\ &\frac{\partial\rho}{\partial t}-\nabla\cdot\left(\rho\nabla
U\right)=0\\\
&{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}U(x,T)=1+\log\frac{\rho(x,T)}{\rho_{ref}(x)},\,\rho(x,0)=\pi(x).}\end{dcases}$
(28)

The main takeaway is that the addition of the quadratic running cost function
produces a generative model that appears to be a time-dependent gradient flow
due to the structure of the continuity equation for $\rho$. Moreover, Onken et
al. (2021) found that OT-flow models are easier and faster to train as it
required fewer parameters. This is likely due to the fact that only a scalar
potential function needs to be learned rather than every component of a
possibly high-dimensional vector field. Moreover, by explicitly constructing a
model informed by the inductive bias of the generative flow, the velocity
field can train faster. This formulation was also explored in Huang et al.
(2022) for discrete normalizing flows, in which the flow map of a continuous
normalizing flow is learned rather than the velocity field.

For completeness, the following optimization problem can be derived from (27),
the log-likelihood formula (22), and the change of variables (21), and is
exactly the OT-flow problem. The objective function can be computed via
training samples from $\pi$:

$\displaystyle\inf_{U}\left\\{\mathbb{E}_{\pi}\left[-\log\rho_{ref}(x(T))-\int_{T}^{0}\Delta
U(x(s),s)\,\text{d}s+\frac{1}{2}\int_{0}^{T}|\nabla
U({x}(s),s)|^{2}\,\text{d}s\right]\right\\}$ (29) $\displaystyle
x(s)=x+\int_{T}^{s}\nabla U(x(t),t)\,\text{d}t,\,x\sim\pi.$

####  3.3.2 Boltzmann generators as MFGs

Boltzmann generators are normalizing flows that have been successfully used
for enhancing sampling from complex distributions that arise in computational
chemistry. Sampling Boltzmann distributions $\pi(x)\propto\exp(-V(x))$ is
challenging as the energy function $V(x)$, while known up to a normalizing
constant, may be highly multimodal. The task is often challenging with
traditional sampling methods such as Markov chain Monte Carlo or molecular
dynamics methods. Boltzmann generators use a combination of generative
modeling and sampling methods to more efficiently explore the energy landscape
Noé et al. (2019); Köhler et al. (2020). Using actual samples from the
distribution $\pi(x)$ produced by some traditional sampling method, a
normalizing flow is constructed, which is then used to construct an importance
sampling distribution to produce more samples of the target distribution.

The optimization problem that is solved to train the importance sampling
distribution can also be cast as a potential MFG, see Section 2.2, as we
demonstrate next. As both the evaluation of target density and samples from it
are accessible, the terminal cost is a convex combination of the forward and
reverse KL divergences. That is, we select

$\displaystyle\mathcal{M}(x,\rho)=\lambda\mathcal{D}_{KL}(\pi\|\rho)+(1-\lambda)\mathcal{D}_{KL}(\rho\|\pi)$
(30)

for some $\lambda\in[0,1]$. In Noé et al. (2019); Köhler et al. (2020),
Boltzmann generators, which are variants of continuous normalizing flows, are
introduced and they correspond to the following MFG, where $L(x,v)=0$ and
$\mathcal{I}(\rho)=0$,

$\displaystyle\min_{v,\rho}\left\\{\lambda\mathcal{D}_{KL}(\pi\|\rho)+(1-\lambda)\mathcal{D}_{KL}(\rho\|\pi):\frac{\partial\rho}{\partial
t}+\nabla\cdot\left(\rho v\right)=0,\,\rho(x,0)=\rho_{ref}(x)\right\\}.$ (31)

The optimality conditions are the same as (25) except for the terminal
condition

$\displaystyle
U(x,T)=\frac{\delta\mathcal{M}}{\delta\rho}(x)=-\lambda\frac{\pi(x)}{\rho(x)}+(1-\lambda)\left(1+\log\frac{\rho(x)}{\pi(x)}\right)$
(32)

Without constraining the feasible velocity fields, the Boltzmann generator as
presented in Noé et al. (2019) is ill-posed in a way identical to the
canonical normalizing flows objective, which we discussed in Section 3.2. In
Köhler et al. (2020), they instead parameterize the normalizing flow velocity
field as the gradient of the potential function, similar to the OT-flow.
However, based on the MFG optimality conditions in (25), we already know that
the velocity field is not of that form, and should instead have a normalized
gradient structure (24). Furthermore, we also know that a gradient structure
for the velocity field is optimal only if there is a quadratic running cost
function as shown in (28). Therefore, to remedy the theoretical inconsistency
in Köhler et al. (2020), we can simply introduce a quadratic running cost.
Based on this MFG-based framework, we will present a corrected version of
Boltzmann generators in Section 7.

##  4 Score-based generative models as solutions to MFGs

We derive the score-based generative modeling (SGM) with stochastic
differential equations framework as presented in Song et al. (2021b) through
mean-field games. We first present an overview of SGM with notation that will
be consistent with the subsequent MFG formulation.

Let $\pi$ be the target data distribution and let
$f:\mathbb{R}^{d}\times\mathbb{R}\to\mathbb{R}^{d}$ be a time-evolving vector
field, and $\sigma:\mathbb{R}\to\mathbb{R}$ be a positive function. Let $Y(t)$
be the diffusion process over $s\in[0,T]$,

$\displaystyle\,\text{d}Y(s)$
$\displaystyle=-f(Y(s),T-s)\,\text{d}t+\sigma(T-s)\,\text{d}W(s)$ (33)
$\displaystyle Y(0)$ $\displaystyle\sim\pi$

which evolves $\eta(\cdot,0)=\pi$ to $\eta(\cdot,T)$. Score-based generative
modeling aims to find a process $X(t)\sim\rho(\cdot,t)$ that reverses the
evolution of $Y(t)$, i.e., $\rho(x,t)=\eta(x,T-t).$ Remarkably, the reverse
process can be expressed in terms of the score function, $\nabla\log\eta(y,s)$
of the diffusion process $Y(s)$ Anderson (1982). For $t\in[0,T]$, $X(t)$
evolves according to

$\displaystyle\,\text{d}X(t)$
$\displaystyle=\left[f(X(t),t)+\sigma(t)^{2}\nabla\log\eta(x,T-t)\right]\,\text{d}t+\sigma(t)\,\text{d}W(t)$
(34) $\displaystyle X(0)$ $\displaystyle\sim\eta(\cdot,T).$

In other words, the reverse process evolves samples from
$\rho(\cdot,0)=\eta(\cdot,T)$ to samples of
$\rho(\cdot,T)=\eta(\cdot,0)=\pi(\cdot)$. In practice, the forward process is
constructed such that $\eta(y,T)$ is close to a normal distribution. For
example, if $f(y,t)\coloneqq y$ and $\sigma(t)\coloneqq\sqrt{2}$ then for
sufficiently large $T$, $\eta(\cdot,T)$ will be approximately normal. New
approximate samples from $\pi$ can then be generated by evolving samples from
the normal distribution through the reverse process.

For generative modeling, we only require learning the score function and then
simulating the reverse SDE. The score function is modeled by some neural
network $\mathsf{s}_{\theta}(y,t)$, which is trained through a score-matching
objective. In essence, $\mathsf{s}_{\theta}(y,t)$ minimizes a least-squares
loss function with respect to the true score function. The explicit score-
matching objective

$\displaystyle
C_{ESM}(\theta)\coloneqq\int_{0}^{T}\frac{1}{2}\mathbb{E}_{\eta(y,s)}\left[\sigma(T-s)^{2}|\mathsf{s}_{\theta}(y,s)-\nabla\log\eta(y,s)|^{2}\right]\,\text{d}s$
(35)

is not computable since $\nabla\log\eta(y,t)$ is not accessible. Instead, the
implicit score-matching objective Song et al. (2020)

$\displaystyle
C_{ISM}(\theta)\coloneqq\int_{0}^{T}\mathbb{E}_{\eta(y,s)}\left[\sigma(T-s)^{2}\left(\frac{1}{2}|\mathsf{s}_{\theta}(y,s)|^{2}+\nabla\cdot\mathsf{s}_{\theta}(y,s)\right)\right]\,\text{d}s$
(36)

and the denoising score-matching objective Vincent (2011); Song and Ermon
(2019), which we omit here, are used. These objective functions avoid the need
to evaluate $\eta(y,s)$ or its derivatives.

###  4.1 Deriving SGM from an MFG

We show that SGM can be derived from a potential MFG (11) with
$L(x,v)=\frac{1}{2}|v|^{2}-\nabla\cdot f$, $\mathcal{I}=0$, and
$\mathcal{M}(\rho)=-\mathbb{E}_{\rho}\left[\log\pi\right]$.

######  Theorem 7.

Let $\pi$ be the target distribution. Score-based generative models are
solutions to the mean-field game

$\displaystyle\inf_{v,\rho}\left\\{-\int_{\mathbb{R}^{d}}\log\pi(x)\rho(x,T)\,\text{d}x+\int_{0}^{T}\int_{\mathbb{R}^{d}}\left(\frac{1}{2}|v(x,t)|^{2}-\nabla\cdot
f(x,t)\right)\rho(x,t)\,\text{d}x\,\text{d}t\right\\}$ (37)
$\displaystyle\text{s.t. }\frac{\partial\rho(x,t)}{\partial
t}+\nabla\cdot\left[(f(x,t)+\sigma(t)v(x,t))\rho(x,t)\right]=\frac{1}{2}\sigma(t)^{2}\Delta\rho(x,t)$
$\displaystyle\rho_{0}(x)=\eta(x,T)$

with controlled dynamics

$\displaystyle\,\text{d}x(t)=\left[f(x(t),t)+\sigma(t)v(x(t),t)\right]\,\text{d}t+\sigma(t)\,\text{d}W(t).$
(38)

The solution of the mean-field game satisfies the optimality conditions

$\displaystyle\begin{dcases}-\frac{\partial U}{\partial t}-f^{\top}\nabla
U+\frac{1}{2}|\sigma\nabla U|^{2}+\nabla\cdot f=\frac{\sigma^{2}}{2}\Delta
U\\\ \frac{\partial\rho}{\partial t}+\nabla\cdot\left((f-\sigma^{2}\nabla
U)\rho\right)=\frac{\sigma^{2}}{2}\Delta\rho\\\
U(x,T)=-\log\pi(x),\,\rho(x,0)=\eta(x,T).\end{dcases}$ (39)

where the optimal velocity field has the representation formula

$\displaystyle v^{*}(x,t)=-\sigma(t)\nabla U(x,t).$ (40)

Moreover, the solution to the HJB equation solves a time-reversed Fokker-
Planck equation

$\displaystyle\frac{\partial\eta(y,s)}{\partial s}$
$\displaystyle=\nabla\cdot(f(y,T-s)\eta(y,s))+\frac{\sigma(T-s)^{2}}{2}\Delta\eta(y,s)$
(41) $\displaystyle\eta(y,0)$ $\displaystyle=\pi(y).$

where

$\displaystyle U(x,t)=-\log\eta(x,T-t),$ (42)

and $\rho^{*}(x,t)=\eta(x,T-t)$. Note that (41) corresponds with the
uncontrolled forward SDE arising in SGM (33), namely

$\displaystyle\,\text{d}y(s)$
$\displaystyle=-f(y(s),T-s)\,\text{d}s+\sigma(T-s)\,\text{d}W_{s}$ (43)
$\displaystyle y(0)$ $\displaystyle\sim\eta(\cdot,0)=\pi(\cdot)\,.$

Finally, the continuity equation in the MFG optimality conditions (39) is the
Fokker-Planck equation corresponding with the reverse SGM dynamics in (34),
i.e. the generative part of the SGM.

Here, the running cost function is a kinetic energy-like term minus the
divergence of $f$, and the terminal cost is the _cross-entropy_ of target
distribution $\pi$ relative to $\rho(\cdot,T)$,
$\mathcal{M}(\rho)=-\mathbb{E}_{\rho}\left[\log\pi\right]$. The cross-entropy
can also be interpreted as the sum of the entropy of $\rho(\cdot,T)$ and the
KL divergence from $\pi$ to $\rho(\cdot,T)$, i.e., observe that
$-\mathbb{E}_{\rho}[\log\pi]=-\mathbb{E}_{\rho}[\log\rho]+\mathcal{D}_{KL}(\rho\|\pi)$.
In the proof of this theorem, we will show that the forward-reverse SDE pair
is _naturally_ induced via the optimality conditions of the MFG.

Proof  First note that the Hamiltonian is

$\displaystyle H(x,p)$ $\displaystyle=\sup_{v}-p^{\top}(f+\sigma
v)-\frac{1}{2}|v|^{2}+\nabla\cdot f$ (44)
$\displaystyle=-f^{\top}p+\frac{1}{2}|\sigma p|^{2}+\nabla\cdot f,$

where the optimizer is $v^{*}=-\sigma p$. The variational derivative of the
cross entropy is

$\displaystyle
M(x,\rho)=-\frac{\delta}{\delta\rho}\mathbb{E}_{\rho}\left[\log\pi\right]=-\log\pi(x).$
(45)

This implies that the optimal velocity field is $v^{*}=-\sigma\nabla U$, where
$U(x,t)$ solves the Hamilton-Jacobi-Bellman equation

$\displaystyle-\frac{\partial U}{\partial t}-f^{\top}\nabla
U+\frac{\sigma^{2}}{2}|\nabla U|^{2}+\nabla\cdot f=\frac{\sigma^{2}}{2}\Delta
U$ (46) $\displaystyle U(x,T)=-\log\pi(x).$

Applying the logarithmic transformation $U(x,t)=-\log\eta(x,T-t)$ Fleming and
Soner (2006), also known as the Cole-Hopf transform Evans (1998); Berner et
al. (2022), to (46), we obtain

$\displaystyle\frac{1}{\eta}\frac{\partial\eta}{\partial
t}+\frac{f^{\top}}{\eta}\nabla\eta+\frac{\sigma^{2}}{2\eta^{2}}|\nabla\eta|^{2}+\nabla\cdot
f=-\frac{\sigma^{2}(\eta\Delta\eta-|\nabla\eta|^{2})}{2\eta^{2}}$
$\displaystyle\implies$ $\displaystyle\frac{\partial\eta}{\partial
t}+f^{\top}\nabla\eta+\frac{\sigma^{2}}{2\eta}|\nabla\eta|^{2}+\eta\nabla\cdot
f=-\frac{\sigma^{2}}{2}\Delta\eta+\frac{\sigma^{2}|\nabla\eta^{2}|}{2\eta}$

which is equivalent to the following linear PDE

$\displaystyle\frac{\partial\eta(x,T-t)}{\partial
t}+\nabla\cdot(f(x,t)\eta(x,T-t))+\frac{\sigma(t)^{2}}{2}\Delta\eta(x,T-t)=0,\,\eta(x,T-T)=\pi(x).$

If we then reparametrize time to be $s=T-t$, the linear PDE is equivalent to
the Fokker-Planck equation

$\displaystyle\frac{\partial\eta(x,s)}{\partial
s}=-\nabla\cdot(-f(x,T-s)\eta(x,s))+\frac{\sigma(T-s)^{2}}{2}\Delta\eta$ (47)
$\displaystyle\eta(x,0)=\pi(x).$

This implies that the optimal velocity
$v^{*}(x,t)=\sigma(t)\nabla\log\eta(x,T-t)$ is related to the solution of the
Fokker-Planck of the uncontrolled SDE

$\displaystyle\,\text{d}y(s)=-f(y(s),T-s)\,\text{d}s+\sigma(T-s)\,\text{d}W(s)$
$\displaystyle y(0)\sim\pi.$

Moreover, we have that

$\rho^{*}(x,t)=\eta(x,T-t)$

is the solution to the Fokker-Planck of the controlled diffusion process in
Theorem 7 for $v^{*}(x,t)=\sigma(t)\nabla\log\eta(x,T-t)$. Observe that

$\displaystyle\frac{\partial\rho^{*}(x,t)}{\partial
t}+\nabla\cdot\left[(f(x,t)+\sigma(t){v^{*}(x,t)})\rho^{*}(x,t)\right]$
$\displaystyle=-\frac{\partial\eta}{\partial
s}+\nabla\cdot\left[(f+\sigma^{2}\nabla\log\eta)\eta\right]$
$\displaystyle=-\frac{\partial\eta}{\partial
s}+\nabla\cdot\left[f\eta+\sigma^{2}\nabla\eta\right]$
$\displaystyle=-\frac{\sigma^{2}}{2}\Delta\eta+\sigma^{2}\Delta\eta$
$\displaystyle=\frac{\sigma^{2}}{2}\Delta\rho^{*}.$

Finally, notice that the initial condition is determined by the solution of
the Fokker-Planck $x(0)\sim\eta(\cdot,T)$.

In Berner et al. (2022), it is shown that the reverse SDE can be derived by
considering a stochastic optimal control problem. However, they pre-supposed
the existence of the forward-reverse SDE pair to derive the optimal control
problem. Through the MFG framework, we show that the forward-reverse SDE pair
is _induced_ through the HJB equation of the MFG optimality condition, which
in this particular setup, is equivalent to the Fokker-Planck equation for
$\eta$ in Theorem 7 via a logarithmic transformation.

######  Remark 8 (Natural boundary conditions of the SGM MFG).

Notice that the initial condition in (37) are implicitly determined by the
optimality conditions of the MFG. This is an example of a _natural boundary
condition_ since the initial density $\rho_{0}(x)$ is not imposed explicitly
in the MFG formulation.

###  4.2 Score-matching objective through duality

We relate the implicit score-matching objective (36) to the mean-field game
formulation we study in Theorem 7. It is unclear that (37) can be related to
(36) directly, mainly due to the fact that in the former, the velocity field
and densities must be optimized simultaneously, while the latter problem has
the density and the velocity field decoupled. To make the connection to score-
matching, we first need to discuss an inherent duality in the SGM MFG.

The proof of Theorem 7 demonstrates that there is a dual nature of the MFG
formulation since the optimality conditions reduce to a pair of Fokker-Planck
equations — one corresponding to the controlled SDE, the other corresponding
to an uncontrolled SDE. The Hamilton-Jacobi-Bellman equation, which comes from
the optimal control problem of the controlled SDE, _induces_ the existence of
the uncontrolled SDE. The duality allows us to _exchange the roles of the
controlled and uncontrolled_ systems through a reparametrization. The
stochastic optimal control objective derived in Berner et al. (2022) makes the
connection to the score-matching objectives through a reparametrization of the
optimal control problem. The discussion we have here will properly justify
that reparametrization considered in Berner et al. (2022); Huang et al. (2021)
using our MFG framework. We first prove a lemma that states the duality of the
MFG and its optimality conditions.

######  Lemma 9.

(Duality and reparametrization) At optimality, the roles of the Fokker-Planck
and HJB equations in Theorem 7 can be exchanged through a reparametrization.
That is, the following two systems are equivalent

$\displaystyle\begin{dcases}-\frac{\partial U(x,t)}{\partial
t}-f(x,t)^{\top}\nabla U(x,t)+\frac{\sigma(t)^{2}}{2}|\nabla
U(x,t)|^{2}+\nabla\cdot f(x,t)=\frac{\sigma(t)^{2}}{2}\Delta U(x,t)\\\
\frac{\partial\rho(x,t)}{\partial
t}+\nabla\cdot\left((f(x,t)-\sigma(t)^{2}\nabla
U(x,t))\rho(x,t)\right)=\frac{\sigma(t)^{2}}{2}\Delta\rho(x,t)\end{dcases}$
(48) $\displaystyle\begin{dcases}\frac{\partial\eta(y,s)}{\partial
s}+\nabla\cdot\left((-g(y,s)-\sigma(T-s)^{2}\nabla
V(y,s))\eta(y,s)\right)=\frac{\sigma(T-s)^{2}}{2}\Delta\eta(y,s)\\\
-\frac{\partial V(y,s)}{\partial s}+g(y,s)^{\top}\nabla
V(y,s)-\frac{\sigma(T-s)^{2}}{2}|\nabla V(y,s)|^{2}-\nabla\cdot
g(y,s)=\frac{\sigma(T-s)^{2}}{2}\Delta V(y,s),\end{dcases}$ (49)

where

$\displaystyle
g(y,s)=f(y,T-s)+\sigma(T-s)v^{*}(y,T-s)=f(y,T-s)-\sigma(T-s)^{2}\nabla
U(y,T-s)\,.$ (50)

Moreover, $U(x,t)=-\log\eta(x,T-t)$, $V(y,s)=-\log\rho(y,T-s)$, and
$\rho(\cdot,t)=\eta(\cdot,T-t)$.

Proof  Since the logarithmic transformation converts the HJB into a time-
reversed Fokker-Planck equation, we can also apply the logarithmic
transformation to the Fokker-Planck equation to produce a HJB equation. From
Theorem 7, if we let $U(x,t)=-\log\eta(x,T-t)$, then the optimality conditions
yield a pair of Fokker-Planck equations,

$\displaystyle\begin{dcases}\frac{\partial\eta(y,s)}{\partial
s}+\nabla\cdot\left((-f(y,T-s))\eta(y,s)\right)=\frac{\sigma(T-s)^{2}}{2}\Delta\eta(y,s)\\\
\frac{\partial\rho(x,t)}{\partial
t}+\nabla\cdot\left((f(x,t)-\sigma(t)^{2}\nabla
U(x,t))\rho(x,t)\right)=\frac{\sigma(t)^{2}}{2}\Delta\rho.\end{dcases}$

To exchange the roles of the uncontrolled FP with the controlled one, define
$g(y,s)=f(y,T-s)-\sigma(T-s)\nabla U(y,T-s)$, and $V(y,s)=U(y,T-s)$. We then
have

$\displaystyle\begin{dcases}\frac{\partial\eta(y,s)}{\partial
s}+\nabla\cdot\left((-g(y,s)-\sigma(T-s)\nabla
V(y,s))\eta(y,s)\right)=\frac{\sigma(T-s)^{2}}{2}\Delta\eta(y,s)\\\
\frac{\partial\rho(x,t)}{\partial
t}+\nabla\cdot\left(g(x,T-t)\rho(x,t)\right)=\frac{\sigma(t)^{2}}{2}\Delta\rho.\end{dcases}$

Lastly, perform a second logarithmic transformation $V(y,s)=-\log\rho(y,T-s)$,
so that we derive

$\displaystyle\begin{dcases}\frac{\partial\eta(y,s)}{\partial
s}+\nabla\cdot\left((-g(y,s)-\sigma(T-s)\nabla
V(y,s))\eta(y,s)\right)=\frac{\sigma(T-s)^{2}}{2}\Delta\eta(y,s)\\\
-\frac{\partial V(y,s)}{\partial s}+g(y,s)^{\top}\nabla
V(y,s)+\frac{\sigma(T-s)^{2}}{2}|\nabla V(y,s)|^{2}-\nabla\cdot
g(y,s)=\frac{\sigma(T-s)^{2}}{2}\Delta V(y,s).\end{dcases}$

Lemma 9 shows two sets of optimal conditions have the same solution through a
reparametrization at the optimum. Therefore, there is a second MFG that
corresponds to the second set of optimality conditions in the lemma which has
the same optimizer as (37). Furthermore, we can show that this second MFG is
equivalent to score-matching.

######  Theorem 10.

Mean-field games (37) and (51)

$\displaystyle\inf_{v,\eta}\left\\{-\int_{\mathbb{R}^{d}}\log\rho(y,0)\eta(y,T)\,\text{d}y+\int_{0}^{T}\int_{\mathbb{R}^{d}}\left(\frac{1}{2}|v(y,T-s)|^{2}+\nabla\cdot
g(y,s)\right)\eta(y,t)\,\text{d}y\,\text{d}t\right\\}$ (51)
$\displaystyle\text{s.t.}\,\frac{\partial\eta(y,s)}{\partial
s}+\nabla\cdot\left[(-g(y,s)+\sigma(T-s)v(y,T-s))\eta(y,s)\right]=\frac{1}{2}\sigma(T-s)^{2}\Delta\eta(y,s)$
$\displaystyle\eta(y,0)=\pi(y),$

with $g=g(y,s)$ defined in Lemma 9, and the implicit score-matching problem
(52)  with $\eta=\eta(y,s)$ given by (41),

$\displaystyle\inf_{\mathsf{s}}\left\\{\int_{0}^{T}\int_{\mathbb{R}^{d}}\sigma(T-s)^{2}\left(\frac{1}{2}|\mathsf{s}(y,s)|^{2}+\nabla\cdot\mathsf{s}(y,s)\right)\eta(y,s)\,\text{d}y\,\text{d}s\right\\}$
(52)

have the same minimizers. Note that the latter functional (52) is identical to
(36).

Proof  The optimality conditions of mean-field games (37) and (51) were shown
to be equivalent in Lemma 9. To show equivalence of MFG (51) to the score-
matching objective, recall from Lemma 9 that at optimality,
$g(y,s)=f(y,T-s)+\sigma(T-s)v^{*}(y,T-s)$, and so

$\displaystyle-\int_{\mathbb{R}^{d}}\log\rho(y,0)\eta^{*}(y,T)\,\text{d}y+$
$\displaystyle\int_{0}^{T}\int_{\mathbb{R}^{d}}\left(\frac{1}{2}|v^{*}(y,T-s)|^{2}+\nabla\cdot(f(y,T-s)+\sigma(T-s)v^{*}(y,T-s))\right)\eta^{*}(y,s)\,\text{d}y\,\text{d}s$
$\displaystyle=\inf_{v}\left\\{-\int_{\mathbb{R}^{d}}\log\rho(y,0)\eta^{*}(y,T)+\int_{0}^{T}\int_{\mathbb{R}^{d}}\left(\frac{1}{2}|v|^{2}+\nabla\cdot(f+\sigma
v)\right)\eta^{*}(y,s)\,\text{d}y\,\text{d}s\right\\},$

where $\eta^{*}(y,s)$ is the optimal density evolution. By Theorem 7 and Lemma
9, (and by some abuse of notation), we know that $\eta^{*}(y,s)=\eta(y,s)$
given by (41). Defining $\sigma(T-s)\mathsf{s}(y,s)=v(y,T-s)$, and ignoring
terms that are independent of $v$, the above problem is equivalent to

$\displaystyle\inf_{\mathsf{s}}\left\\{\int_{0}^{T}\int_{\mathbb{R}^{d}}\sigma(T-s)^{2}\left(\frac{1}{2}|\mathsf{s}(y,s)|^{2}+\nabla\cdot\mathsf{s}\right)\eta(y,s)\,\text{d}y\,\text{d}s\right\\},$

which is precisely implicit score-matching.

For connections between the score-matching objectives, (35) and (36), and the
evidence lower bound (ELBO), see Huang et al. (2021); Song et al. (2021a). The
equivalence of implicit score matching and the frequently used denoising
score-matching can be found in Berner et al. (2022); Vincent (2011).

###  4.3 MFGs unravel the curious mathematical serendipity of score-based
generative modeling

We take a broader perspective on the curious mathematical structure of the
mean-field game associated with score-based generative modeling. The
optimality conditions of any MFG (11) are two coupled PDEs (12)— a continuity
equation that is solved forwards in time, which determines the evolution of
the density, and a HJB equation that is solved backwards in time, which
determines the optimal control. In general, it is quite difficult to find
structure that simplifies the coupled set of PDEs. It is intriguing, however,
that for the MFG associated with SGM (37), a series of curious mathematical
coincidences conspire to create a mean-field game that not only decouples the
pair of PDEs (39), but also yields two stochastic differential equations – one
controlled (38), one uncontrolled (43), where the optimal control of the
controlled SDE (46) is determined by the density evolution (47) of the
uncontrolled SDE. Furthermore, the duality of the optimality conditions (Lemma
9) shows that there is, in fact, a pair of MFGs, (37) and (51), that share the
same solution, one of which (51) can be related to implicit score-matching
(36).

We emphasize three key characteristics of the MFG formulation of SGM that
makes the structure described above possible:

  * •

The logarithmic transformation that converts the HJB equation (46) into a
Fokker-Planck (47) is only possible if the controlled dynamical system (38) is
a stochastic differential equation. The presence of the Laplacian operator on
the right hand side of (46) is crucial for the logarithmic transformation to
convert the nonlinear Hamilton-Jacobi equation into a linear PDE.

  * •

Both terms in the running cost $L(x,v)=|v|^{2}/2-\nabla\cdot f$ are important.
The quadratic cost function yields a gradient structure for the velocity
field, which we later find is a score function. The divergence term
$\nabla\cdot f$ is needed so that after a logarithmic transformation is
applied to the HJB, the advection term of a Fokker-Planck equation (47)
appears.

  * •

The terminal cost is the _cross-entropy_ of $\pi$ with respect to the terminal
density $\rho(\cdot,T)$. This yields a terminal condition in the HJB equations
(46) that is independent of the terminal density of the Fokker-Planck.
Therefore, the mean-field game (37) is, in fact, a stochastic optimal control
problem, _cf._ the terminal conditions of the normalizing flow optimality
conditions in (19). The cross-entropy cost also results in the solution to the
HJB to be the logarithm of a density function. Further connections between the
score-matching objective and the cross-entropy is discussed in Vahdat et al.
(2021).

The first two features imply that the HJB can be transformed into a time-
reversed Fokker-Planck equation. These two characteristics alone do not
necessarily imply that the HJB corresponds to the density evolution of an
uncontrolled SDE. The third point regarding the cross-entropy is crucial since
it determines that the terminal condition of the time-reversed Fokker-Planck
equation is, indeed, a density function. Altogether, these three
characteristics imply that the solution to the HJB (46) is related to the
solution of an uncontrolled SDE. Without the cross-entropy terminal cost in
SGM, the pair of SDEs still exists, except that both of them are controlled
SDEs. In fact, the problem is the Schrödinger bridge problem, which we discuss
in Section 4.5.

A fourth fact allows us to relate the MFG objective function to the implicit
score-matching objective:

  * •

The dual nature of the pair of Fokker-Planck equations imply that there are
two equivalent mean-field games with the same solution (Lemma 9). That is,
there is no requirement that one of the processes must be controlled or
uncontrolled. A simple reparametrization (50) can swap the roles of the two
SDEs.

######  Remark 11 (A shooting method analogy).

Score-based generative modeling has been found to be superior to normalizing
flows Song et al. (2021b), and the MFG perspective can give some intuitive
insight into why this may be. One can interpret generative modeling as finding
the trajectory on the space of measures that connects an initial distribution
$\rho_{0}$ to the target distribution $\pi$. The MFG analysis of SGM shows
that because the pair of PDEs yields a pair of SDEs, (38) and (43), that
invert each other, the trajectory on the space of measures is prescribed _a
priori_. Therefore, score-matching is simply finding a least squares
approximation (35) to that trajectory.

In contrast, normalizing flows based on minimizing the KL divergence (17) is
necessarily a mean-field game, and so the _velocity and density have to be
optimized simultaneously_. Unlike SGM, there is no _elegant relation between
the density and the optimal velocity field_. Intuitively, and in practice, the
CNF problem can be thought of as a shooting problem on the space of
probability measures. We first guess a velocity field, which results in an
approximate density to $\pi$. Based on the discrepancy between the approximate
and target distribution, the velocity field is updated so that the approximate
density is closer to $\pi$. Reasoning by analogy, two point boundary value
problems for ODEs on $\mathbb{R}^{d}$ are generally numerically difficult to
solve with the shooting method. Therefore, we argue by analogy that
normalizing flows are more difficult to train than SGMs.

######  Remark 12 (Score-based generative models as normalizing flows).

By comparing Theorem 2 with Theorem 7, we can interpret score-based generative
models as a normalizing flow with initial distribution $\eta(x,T)$ defined in
(41). Specifically, SGMs are a normalizing flow with stochastic dynamics, a
quadratic running cost, and the cross-entropy as the terminal cost.

###  4.4 Score-based probability flows as solutions of MFGs

The probability flow ODE presented in Maoutsa et al. (2020); Song et al.
(2021b) defines a system of interacting particles with deterministic dynamics
that approximates the solution of the Fokker-Planck equation. It also provides
a connection between normalizing flows and score-based generative models. We
present a theorem where we show that the probability flow can derived via a
mean-field game without appealing to its relation to score-based models.

We show that the probability flow ODE can be derived from a potential MFG with
$L(x,v)=\frac{1}{2}|v|^{2}-\frac{1}{2}\nabla\cdot f$,
$\mathcal{I}(\rho)=\mathbb{E}_{\rho}\left[\frac{\sigma^{2}}{8}|\nabla\log\rho|^{2}\right]$,
and $\mathcal{M}(\rho)=-\mathbb{E}_{\rho}[\log\pi]$.

######  Theorem 13.

Let $\pi$ be the target distribution. The score-based probability flow ODE is
the solution of the mean-field game

$\displaystyle\inf_{v,\rho}\left\\{-\frac{1}{2}\mathbb{E}_{\rho(x,T)}\left[\log\pi(x)\right]+\int_{0}^{T}\mathbb{E}_{\rho(x,t)}\left[\left(\frac{\sigma(t)^{2}}{8}|\nabla\log\rho(x,t)|^{2}+\frac{|v(x,t)|^{2}-\nabla\cdot
f(x,t)}{2}\right)\right]\,\text{d}t\right\\}$ (53) $\displaystyle\text{s.t.
}\frac{\partial\rho(x,t)}{\partial
t}+\nabla\cdot\left[(f(x,t)+\sigma(t)v(x,t))\rho(x,t)\right]=0$
$\displaystyle\rho_{0}(x)=\eta(x,T)$

with deterministic dynamics

$\displaystyle\,\text{d}x(t)=\left[f(x(t),t)+\sigma(t)v(x(t),t)\right]\,\text{d}t.$
(54)

The solution of the mean-field game satisfies the optimality conditions

$\displaystyle\begin{dcases}-\frac{\partial U}{\partial t}-f^{\top}\nabla
U+\frac{1}{2}|\sigma\nabla U|^{2}+\frac{1}{2}\nabla\cdot
f=\frac{\sigma^{2}|\nabla\rho|^{2}}{8\rho^{2}}-\frac{\sigma^{2}\Delta\rho}{4\rho}\\\
\frac{\partial\rho}{\partial t}+\nabla\cdot\left((f-\sigma^{2}\nabla
U)\rho\right)=0\\\
U(x,T)=-\frac{1}{2}\log\pi(x),\,\rho(x,0)=\eta(x,T).\end{dcases}$ (55)

where the optimal velocity field has the representation formula

$\displaystyle v^{*}(x,t)=-\sigma(t)\nabla U(x,t).$ (56)

Moreover, the solution to the HJB equation solves a time-reversed Fokker-
Planck equation

$\displaystyle\frac{\partial\eta(y,s)}{\partial s}$
$\displaystyle=\nabla\cdot(f(y,T-s)\eta(y,s))+\frac{\sigma(T-s)^{2}}{2}\Delta\eta(y,s)$
(57) $\displaystyle\eta(y,0)$ $\displaystyle=\pi(y).$

where

$\displaystyle U(x,t)=-\frac{1}{2}\log\eta(x,T-t),$ (58)

and $\rho^{*}(x,t)=\eta(x,T-t)$. Again, (57) corresponds with the uncontrolled
SDE:

$\displaystyle\,\text{d}y(s)$
$\displaystyle=-f(y(s),T-s)\,\text{d}s+\sigma(T-s)\,\text{d}W_{s}$ (59)
$\displaystyle y(0)$ $\displaystyle\sim\eta(\cdot,0)=\pi(\cdot)\,.$

Proof  First note that the Hamiltonian is

$\displaystyle H(x,p)$ $\displaystyle=\sup_{v}-p^{\top}(f+\sigma
v)-\frac{1}{2}|v|^{2}+\frac{1}{2}\nabla\cdot f$
$\displaystyle=-p^{\top}f+\frac{\sigma^{2}}{2}|p|^{2}+\frac{1}{2}\nabla\cdot
f,$ (60)

with the optimizer being $v^{*}=-\sigma p$. Next, notice that the interaction
cost functional is

$\displaystyle\mathcal{I}(\rho)=\int_{\mathbb{R}^{d}}\frac{\sigma^{2}}{8}|\nabla\log\rho(x,t)|^{2}\rho(x,t)\,\text{d}x.$
(61)

We compute the variational derivative; for some test function $\chi$, such
that $\int\chi\,\text{d}x=0$, observe that

$\displaystyle\lim_{\epsilon\to
0}\frac{\mathcal{I}(\rho+\epsilon\chi)-\mathcal{I}(\rho)}{\epsilon}$
$\displaystyle=\lim_{\epsilon\to
0}\frac{\sigma^{2}}{8\epsilon}\int_{\mathbb{R}^{d}}\frac{|\nabla(\rho+\epsilon\chi)|^{2}}{\rho+\epsilon\chi}-\frac{|\nabla\rho|^{2}}{\rho}\,\text{d}x$
$\displaystyle=\lim_{\epsilon\to
0}\frac{\sigma^{2}}{8\epsilon}\int_{\mathbb{R}^{d}}\frac{\rho|\nabla\rho|^{2}+2\rho\epsilon\nabla\rho\cdot\nabla\chi+\epsilon^{2}\rho|\nabla\chi|^{2}-\rho|\nabla\rho|^{2}-\epsilon\chi|\nabla\rho|^{2}}{\rho(\rho+\epsilon\chi)}\,\text{d}x$
$\displaystyle=\lim_{\epsilon\to
0}\frac{\sigma^{2}}{8}\int_{\mathbb{R}^{d}}\frac{2\rho\nabla\rho\cdot\nabla\chi+\epsilon\rho|\nabla\chi|^{2}-\chi|\nabla\rho|^{2}}{\rho(\rho+\epsilon\chi)}\,\text{d}x$
$\displaystyle=\frac{\sigma^{2}}{8}\int_{\mathbb{R}^{d}}2\nabla\log\rho\cdot\nabla\chi-\frac{|\nabla\rho|^{2}}{\rho^{2}}\chi\,\text{d}x$
$\displaystyle=\int_{\mathbb{R}^{d}}\left(-\frac{\sigma^{2}}{4}\Delta\log\rho-\frac{\sigma^{2}|\nabla\rho|^{2}}{8\rho^{2}}\right)\chi\,\text{d}x$
$\displaystyle=\int_{\mathbb{R}^{d}}\left(-\frac{\sigma^{2}\Delta\rho}{4\rho}+\frac{\sigma^{2}|\nabla\rho|^{2}}{4\rho^{2}}-\frac{\sigma^{2}|\nabla\rho|^{2}}{8\rho^{2}}\right)\chi\,\text{d}x$
$\displaystyle=\int_{\mathbb{R}^{d}}\left(-\frac{\sigma^{2}\Delta\rho}{4\rho}+\frac{\sigma^{2}|\nabla\rho|^{2}}{8\rho^{2}}\right)\chi\,\text{d}x.$

Therefore, the HJB-FP system of PDEs is

$\displaystyle\begin{dcases}-\frac{\partial U}{\partial t}-f^{\top}\nabla
U+\frac{1}{2}\sigma^{2}|\nabla U|^{2}+\frac{1}{2}\nabla\cdot
f=\frac{\sigma^{2}|\nabla\rho|^{2}}{8\rho^{2}}-\frac{\sigma^{2}\Delta\rho}{4\rho}\\\
\frac{\partial\rho}{\partial t}+\nabla\cdot(\rho(f-\sigma^{2}\nabla U))=0\\\
\rho(x,0)=\eta(x,T),\,U(x,T)=-\frac{1}{2}\log\pi(x),\end{dcases}$ (62)

and by the derivation of the Hamiltonian, the optimal velocity is
$v^{*}(x,t)=-\sigma(t)\nabla U(x,t)$. We claim that the solution of the HJB-FP
system is of the form $U(x,t)=-\frac{1}{2}\log\eta(x,T-t)$ and
$\rho(x,t)=\eta(x,T-t)$, and that $\eta$ solves a Fokker-Planck equation. To
see that the optimal $U$ and $\rho$ can be expressed in terms of the density
of an uncontrolled SDE, note the following computations. Define $s=T-t$, and
observe that

$\displaystyle-\frac{1}{2\eta}\frac{\partial\eta}{\partial
s}+\frac{1}{2}f^{\top}\nabla\log\eta+\frac{\sigma^{2}}{8}|\nabla\log\eta|^{2}+\frac{1}{2}\nabla\cdot
f=\frac{\sigma^{2}|\nabla\eta|^{2}}{8\eta^{2}}-\frac{\sigma^{2}\Delta\eta}{4\eta}$
$\displaystyle\implies$ $\displaystyle\frac{\partial\eta}{\partial
s}-f^{\top}\nabla\eta-\frac{\sigma^{2}}{4}\eta|\nabla\log\eta|^{2}-\eta\nabla\cdot
f=\frac{\sigma^{2}\Delta\eta}{2}-\frac{\sigma^{2}}{4}\eta|\nabla\log\eta|^{2}$
$\displaystyle\implies$ $\displaystyle\frac{\partial\eta(x,s)}{\partial
s}+\nabla\cdot\left(-f(x,T-s)\eta(x,s)\right)=\frac{\sigma(T-s)^{2}}{2}\Delta\eta(x,s).$

With the time and logarithmic transformation, the terminal condition becomes
an initial condition, where $\eta(x,0)=\pi(x)$. The HJB, once again, becomes
the Fokker-Planck of an uncontrolled SDE.

Lastly, we check that $\rho(x,t)=\eta(x,T-t)$ is the solution of the
controlled SDE with the optimal control $v^{*}(x,t)=-\sigma(t)\nabla U(x,t)$:

$\displaystyle\frac{\partial\rho}{\partial
t}+\nabla\cdot(\rho(f-\sigma^{2}\nabla U))$
$\displaystyle=-\frac{\partial\eta}{\partial
s}+\nabla\cdot\left[\eta\left(f+\frac{\sigma^{2}}{2}\nabla\log\eta\right)\right]$
$\displaystyle=-\frac{\partial\eta}{\partial
s}+\nabla\cdot(f\eta)+\frac{\sigma^{2}}{2}\Delta\eta$ $\displaystyle=0.$

The interaction term here is the Fisher information functional of $\rho$. The
MFG corresponding to the score probability flow in Equation (53) can also be
interpreted as a regularization of the Benamou-Brenier problem. See Chen et
al. (2021); Pavon et al. (2021) for further discussion about the Fisher
information functional, Schrödinger bridges, and the Benamou-Brenier formula.

###  4.5 Schrödinger bridges

We discuss connections between the Schrödinger bridge problem (SBP), score-
based generative modeling, and mean-field games. While there are many
formulations and variations of the SBP, see Chen et al. (2021); Pavon et al.
(2021); Léonard (2014), we focus on the optimal control version. Let
$\rho_{0}$ and $\rho_{T}=\pi$ be two distributions. The Schrödinger bridge
problem aims to find a stochastic process whose $t=0$ and $t=T$ marginals are
$\rho_{0}$ and $\rho_{T}$ respectively,

$\displaystyle\inf_{v}\left\\{\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)\,\text{d}x\,\text{d}t\right\\}$
(63) $\displaystyle\text{s.t. }\,\text{d}x(t)=\sigma
v(x(t),t)\,\text{d}t+\sigma\,\text{d}W(t)$ $\displaystyle
x(0)\sim\rho_{0},\,x(T)\sim\rho_{T}=\pi.$

Here $\sigma>0$ is a constant. From Léonard (2014); Pavon et al. (2021), it is
shown that the solution to the SBP is the control

$\displaystyle v^{*}(x,t)=\sigma\nabla\log\phi(x,t)$ (64)

where $\phi(x,t)$ is the solution of the so-called Schrödinger system Chen et
al. (2021)

$\displaystyle\begin{dcases}\frac{\partial\phi}{\partial
t}+\frac{\sigma^{2}}{2}\Delta\phi=0\\\ \frac{\partial\widehat{\phi}}{\partial
t}-\frac{\sigma^{2}}{2}\Delta{\widehat{\phi}}=0\\\
\phi(x,0)\widehat{\phi}(x,0)=\rho_{0}(x),\,\phi(x,T)\widehat{\phi}(x,T)=\pi(x).\end{dcases}$
(65)

The above stochastic optimal control problem (63) has a form that is quite
similar to a mean-field game. The main difference is that the terminal cost is
instead a terminal constraint $\rho_{T}=\pi$. It is possible to derive the
Schrödinger system (65) by interpreting the SBP as a mean-field game. To do
so, we first consider a more general MFG.

####  4.5.1 MFGs, Schrödinger bridges, and a system of forward-backward heat
equations

Consider the mean-field game with a quadratic running cost
$L(x,v)=\frac{1}{2}|v|^{2}$ and arbitrary interaction and terminal cost
functionals $\mathcal{I}$ and $\mathcal{M}$. Let $\sigma$ be constant. We have

$\displaystyle\inf_{v,\rho}\left\\{\mathcal{M}(\rho(\cdot,T))+\int_{0}^{T}\mathcal{I}(\rho(\cdot,t))\,\text{d}t+\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)\,\text{d}x\,\text{d}t\right\\}$
(66) $\displaystyle\text{s.t. }\frac{\partial\rho}{\partial
t}+\nabla\cdot(v\rho)=\frac{\sigma^{2}}{2}\Delta\rho,\,\rho(x,0)=\rho_{0}.$

This particular mean-field game has been previously studied due to the
structure of the resulting optimality conditions Guéant (2012). One can easily
check that the Hamiltonian is quadratic in the co-state
$H(x,p)=\frac{1}{2}|p|^{2}$, and the optimality conditions are

$\displaystyle\begin{dcases}-\frac{\partial U}{\partial t}+\frac{1}{2}|\nabla
U|^{2}-\frac{\sigma^{2}}{2}\Delta U=I(x,\rho)\\\ \frac{\partial\rho}{\partial
t}-\nabla\cdot(\rho\nabla U)-\frac{\sigma^{2}}{2}\Delta U=0\\\
\rho(x,0)=\rho_{0}(x),\,U(x,T)=M(x,\rho(x,T))\\\ \end{dcases}$ (67)

In Guéant (2012), a variable transformation, which can be interpreted as a
generalized logarithmic transformation, is used to reduce (67) into a system
of two inhomogeneous heat equations with related source terms. Define
$\psi(x,t)=\exp\left(-\frac{U(x,t)}{\sigma^{2}}\right)$ and
$\widehat{\psi}(x,t)=\rho(x,t)\exp\left(\frac{U(x,t)}{\sigma^{2}}\right)$. The
optimality conditions (67) reduce to

$\displaystyle\begin{dcases}\frac{\partial\psi}{\partial
t}+\frac{\sigma^{2}}{2}\Delta\psi=\frac{1}{\sigma^{2}}I(x,\psi\widehat{\psi})\psi\\\
\frac{\partial\widehat{\psi}}{\partial
t}-\frac{\sigma^{2}}{2}\Delta\widehat{\psi}=-\frac{1}{\sigma^{2}}I(x,\psi\widehat{\psi})\widehat{\psi}\\\
\psi(x,T)=\exp\left(-\frac{M(x,\rho(x,T))}{\sigma^{2}}\right)\\\
\widehat{\psi}(x,0)=\frac{\rho(x,0)}{\psi(x,0)}.\end{dcases}$ (68)

With the variable transformation, it is simple to see that this system closely
matches the Schrödinger system (65). The main differences are that there is no
interaction term $I$, and that the terminal condition is not a strict terminal
constraint. Replacing a terminal constraint that is difficult to enforce with
a penalty is sometimes considered to relax the optimization problem Genevay et
al. (2017).

The link between the SBP and mean-field games has attracted recent attention.
For example, Liu et al. (2022) argued the connection via the MFG objective, as
we do here, and they further generalize the SBP by introducing an interaction
cost function, which naturally arises in the MFG formulation, by choosing $I$
to be nonzero. Moreover, using solutions of the SBP for generative modeling
has also been an active area of research. For example, Wang et al. (2021a); De
Bortoli et al. (2021); Shi et al. (2023) have trained generative models based
on solutions of the Schrödinger bridge problem. Moreover, Chen et al. (2022)
learns Schrödinger bridge-based generative models using forward-backward SDE
theory.

##  5 Generative models based on Wasserstein gradient flows as solutions to
MFGs

In this section we discuss how the Wasserstein gradient flow can be derived
with mean-field games. Specifically, we construct a sequence of mean-field
games whose limit is the Wasserstein gradient flow. This connection will allow
us to introduce relaxations to the Wasserstein gradient flows and may lead to
future computational investigations.

###  5.1 Wasserstein gradient flows as solutions to MFG

We now formulate the Wasserstein gradient flow as the solution of an MFG. The
key results that allow us to connect these two concepts come from the study of
general metric gradient flows through variational principles. The main idea
here is that a variational principle appears by combining gradient flows with
elliptic regularization. It was shown in Rossi et al. (2011); Segatti (2013);
Rossi et al. (2019); Mielke and Stefanelli (2011) that gradient flows on
metric spaces can be described as minimizers of a so-called weighted energy
dissipation (WED) functional. For Wasserstein space, this minimization problem
can be related to the mean-field game described in (69). Let
$\mathcal{F}(\nu)$ be a geodesically convex energy functional on
$\mathcal{P}(\Omega)$.111A constant speed geodesic between two measures
$\mu,\nu\in\mathcal{P}(\Omega)$ is defined as
$\mu_{t}=((1-t)\text{Id}+t\mathcal{T})_{\sharp}\mu$, where
$\mathcal{T}_{\sharp}\mu=\nu$, $\mathcal{T}$ is the optimal map, and
$t\in[0,1]$. A functional $\mathcal{F}(\rho)$ is $\lambda$-geodesically convex
if $\mathcal{F}(\mu_{t})\leq(1-t)\mathcal{F}(\mu)+t\mathcal{F}(\nu)$.  The
following theorem establishes the connection between Wasserstein gradient
flows and MFGs:

######  Theorem 14.

For any $\epsilon>0$ and geodesically convex functional $\mathcal{F}(\rho)$,
the MFG

$\displaystyle\inf_{v,\rho}\left\\{\mathcal{F}(\rho(\cdot,T))e^{-T/\epsilon}+\int_{0}^{T}\frac{e^{-t/\epsilon}}{\epsilon}\mathcal{F}(\rho(\cdot,t))\,\text{d}t+\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{e^{-t/\epsilon}}{2}|v(x,t)|^{2}\rho(x,t)\,\text{d}x\,\text{d}t\right\\}$
(69) $\displaystyle\text{s.t. }\frac{\partial\rho}{\partial
t}+\nabla\cdot(v\rho)=0$ $\displaystyle\rho(x,0)=\rho_{0}(x)$

with dynamics

$\displaystyle\frac{\,\text{d}x}{\,\text{d}t}=v(x(t),t)$ (70)

satisfies the optimality conditions

$\displaystyle\begin{dcases}-\epsilon\frac{\partial U}{\partial
t}+U+\frac{\epsilon}{2}|\nabla U|^{2}=\frac{\delta\mathcal{F}}{\delta\rho}\\\
\frac{\partial\rho}{\partial t}-\nabla\cdot(\rho\nabla U)=0\\\
U(x,T)=\frac{\delta\mathcal{F}}{\delta\rho}(\rho(\cdot,T)),\,\rho(x,0)=\rho_{0}(x),\end{dcases}$
(71)

with optimal velocity field $v^{*}(x,t)=-\nabla U(x,t)$. In particular, as
$\epsilon\to 0$, the solution is the Wasserstein gradient flow, in which

$\displaystyle\begin{dcases}U=\frac{\delta\mathcal{F}}{\delta\rho}\\\
\frac{\partial\rho}{\partial
t}=\nabla\cdot\left(\rho\nabla\frac{\delta\mathcal{F}}{\delta\rho}\right)\\\
\rho(x,0)=\rho_{0}(x).\end{dcases}$ (72)

Proof

To derive the optimality conditions, note that the Hamiltonian is

$\displaystyle H(x,p,t)$
$\displaystyle=\sup_{v}\left[-p^{\top}v-\frac{e^{-t/\epsilon}}{2}|v|^{2}\right]$
$\displaystyle=\frac{e^{t/\epsilon}}{2}|p|^{2},$ (73)

where the maximum is attained when $v^{*}=-e^{t/\epsilon}p$. Therefore, the
optimality conditions are

$\displaystyle\begin{dcases}-\frac{\partial U_{\epsilon}}{\partial
t}+\frac{e^{t/\epsilon}}{2}|\nabla
U_{\epsilon}|^{2}=\frac{e^{-t/\epsilon}}{\epsilon}\frac{\delta\mathcal{F}}{\delta\rho}\\\
\frac{\partial\rho}{\partial t}-\nabla\cdot\left(\rho e^{t/\epsilon}\nabla
U_{\epsilon}\right)=0\\\
U_{\epsilon}(x,T)=e^{-T/\epsilon}\frac{\delta\mathcal{F}}{\delta\rho}(\rho(\cdot,T)),\,\rho(x,0)=\rho_{0}(x).\end{dcases}$
(74)

We may simplify the optimality conditions further. Let
$U^{\epsilon}(x,t)=U(x,t)e^{-t/\epsilon}$, then the HJB equation is

$\displaystyle-\frac{\partial U}{\partial
t}e^{-t/\epsilon}+\frac{e^{-t/\epsilon}}{\epsilon}U+\frac{e^{-t/\epsilon}}{2}|\nabla
U|^{2}=\frac{e^{-t/\epsilon}}{\epsilon}\frac{\delta\mathcal{F}}{\delta\rho},$

which after simplifying, we get the desired result

$\displaystyle\begin{dcases}-\epsilon\frac{\partial U}{\partial
t}+U+\frac{\epsilon}{2}|\nabla U|^{2}=\frac{\delta\mathcal{F}}{\delta\rho}\\\
\frac{\partial\rho}{\partial t}-\nabla\cdot\left(\rho\nabla U\right)=0.\\\
U(x,T)=\frac{\delta\mathcal{F}}{\delta\rho}(\rho(\cdot,T)),\,\rho(x,0)=\rho_{0}(x).\end{dcases}$
(75)

Moreover, sending $\epsilon\to 0$, and assuming smoothness of solutions, we
obtain the Wasserstein gradient flow. This limit is well-defined based on the
results in Rossi et al. (2011, 2019).

######  Remark 15 (Regularized Wasserstein gradient flow).

Alternatively, we may rewrite the continuity equation in (75) as a regularized
Wasserstein gradient flow:

$\displaystyle\frac{\partial\rho}{\partial
t}=\nabla\cdot\left(\rho\nabla\frac{\delta\mathcal{F}}{\delta\rho}\right)+\epsilon\nabla\cdot\left(\rho\nabla\left(\frac{\partial
U}{\partial t}-\frac{1}{2}|\nabla U|^{2}\right)\right)$ (76)

Again, we observe that if $\epsilon\to 0$, we recover the Wasserstein gradient
flow.

Finally, we note that in the $\epsilon\to\infty$ limit, the MFG (69) converges
to a regularized optimal transport problem, as discussed in detail in Section
7.4. As a result, we can interpret different choices of $\epsilon$ as
_interpolating_ between Wasserstein gradient flows ($\epsilon\to 0$) and
(regularized) optimal transport ($\epsilon\to\infty$).

###  5.2 Mean-field games with relaxation

An alternative asymptotic analysis for $\epsilon\to 0$ can be carried out by
rewriting the MFG system (75) in relaxation form, Jin and Xin (1995), namely
as

$\displaystyle\begin{dcases}\frac{\partial U}{\partial t}-\frac{1}{2}|\nabla
U|^{2}=\frac{1}{\epsilon}\left(U-\frac{\delta\mathcal{F}}{\delta\rho}\right)\\\
\frac{\partial\rho}{\partial t}-\nabla\cdot\left(\rho\nabla U\right)=0\\\
U(x,T)=\frac{\delta\mathcal{F}}{\delta\rho}(\rho(\cdot,T)),\,\rho(x,0)=\rho_{0}(x).\end{dcases}$
(77)

We refer to (77) as an MFG with relaxation due to the relaxation term
$\frac{1}{\epsilon}\left(U-\frac{\delta\mathcal{F}}{\delta\rho}\right)$ in the
Hamilton-Jacobi equation. We note the resemblance in structure to the
relaxation approximations for hyperbolic conservation laws introduced in Jin
and Xin (1995). In the latter work, an asymptotic analysis based on a Chapman-
Enskog expansion demonstrates that the relaxation approximations converge as
$\epsilon\to 0$ to the (physically correct) entropy solutions of multi-
dimensional hyperbolic conservation laws. Here the Chapman-Enskog expansion of
(77) is due to the relaxation term
$\frac{1}{\epsilon}\left(U-\frac{\delta\mathcal{F}}{\delta\rho}\right)$ which
in the $\epsilon\ll 1$ regime, enforces a local equilibrium — hence a Chapman-
Enskog expansion — which takes the form

$U(x,t)=\frac{\delta\mathcal{F}}{\delta\rho}+\epsilon v(x,t)\,.$ (78)

We can substitute this in (77) to obtain

$\displaystyle\begin{dcases}\frac{\partial U}{\partial t}-\frac{1}{2}|\nabla
U|^{2}=v\\\ \frac{\partial\rho}{\partial
t}=\nabla\cdot\left(\rho\nabla\frac{\delta\mathcal{F}}{\delta\rho}\right)+\epsilon\nabla\cdot\left(\rho\nabla
v\right)\\\
U(x,T)=\frac{\delta\mathcal{F}}{\delta\rho}(\rho(\cdot,T)),\,\rho(x,0)=\rho_{0}(x).\end{dcases}$
(79)

As $\epsilon\to 0$ the continuity equation yields the Wasserstein gradient
flow

$\displaystyle\frac{\partial\rho}{\partial
t}=\nabla\cdot\left(\rho\nabla\frac{\delta\mathcal{F}}{\delta\rho}\right)\,,\quad\rho(x,0)=\rho_{0}(x).$
(80)

In the derivations of Jin and Xin (1995) the authors also had to resolve the
next order term $\mathcal{O}(\epsilon)$ to ensure the entropy condition is
satisfied for the hyperbolic conservation laws. In our case, the zeroth order
term in the continuity equation in (79) is the Wasserstein gradient flow which
is always well-posed. Therefore, resolving the next order term
$\epsilon\nabla\cdot\left(\rho\nabla v\right)$, while feasible, appears to be
less essential.

######  Remark 16 (Well-posedness and boundary layers).

An important use of the relaxation formulation (77) is related to the well-
posedness of the MFG formulation. Indeed, the corresponding Chapman-Enskog
expansion (78) allows us to justify the choice of the terminal condition
$\mathcal{M}=\mathcal{F}$ in the potential MFG in Theorem 14, as follows. In
principle, the choice of any $\mathcal{M}$ instead of $\mathcal{F}$ in (69),
may appear to be possible or even user-defined. However in Theorem 14 we pick
$\mathcal{M}=\mathcal{F}$ to avoid the boundary layer in the Hamilton Jacobi
equation at $t=T$ and thus facilitate the $\epsilon\to 0$ limit, by ensuring
that $U=\frac{\delta\mathcal{F}}{\delta\rho}$ at $t=T$. By a boundary layer we
mean that a choice $\mathcal{M}\neq\mathcal{F}$ would imply a mismatch between
the terminal condition in (77), which in this case would be
$U(x,T)=\frac{\delta\mathcal{M}}{\delta\rho}(\rho(\cdot,T))$, and the Chapman-
Enskog expansion (78), as $\epsilon\ll 1$. In this sense, the Chapman-Enskog
expansion (78) of the relaxation MFG (77) demonstrates that the choice
$\mathcal{M}=\mathcal{F}$ constitutes a “natural” boundary condition for the
MFG (69).

###  5.3 Applications to generative modeling

Wasserstein gradient flows have applications in many problems in machine
learning, including generative modeling and sampling. A classic example of a
Wasserstein gradient flow is the overdamped Langevin dynamics, where the
interaction functional is the KL divergence. This celebrated result was first
noted in Jordan et al. (1998). Using this gradient flow structure for sampling
and inference has been explored in Maoutsa et al. (2020); Garbuno-Inigo et al.
(2020); Wang and Li (2022); Wang et al. (2021b); Liu and Wang (2016); Nüsken
(2023). Particle-based, non-parametric implicit generative models based on
gradient flows is an active area of research. For example, the use of the
kernelized maximum mean discrepancy for generative modeling was explored in
Arbel et al. (2019). Moreover, discretizations of the gradient flows have been
used to construct implicit generative models Arbel et al. (2019); Glaser et
al. (2021); Gu et al. (2022); Liutkus et al. (2019); Ansari et al. (2020). A
broader class of interaction functionals based on the generalization of
$f$-divergences and integral probability metrics for implicit generative
models in the small data regime was presented in Gu et al. (2022).

There has also been some investigation into the training of generative
adversarial nets (GANs) using the tools of Wasserstein gradient flows. The
training of a GAN can be interpreted as a gradient flow on Wasserstein space,
and implicit generative models can be interpreted as particles moving
according to the gradient of a learned discriminator of a GAN. Studying the
convergence of a restricted class of GANs as a Wasserstein gradient flow was
explored in Mroueh et al. (2019).

##  6 Enhancing generative models with MFGs: HJB regularizers

Characterizing generative flows in terms of MFGs provide formulations of
generative models as solutions to partial differential equations. In
particular, the optimal velocity field that is typically trained by solving an
optimization problem is the solution to the Hamilton-Jacobi-Bellman equation.
The use of HJB equations for accelerating training of generative models has
been studied previously in specific flows Onken et al. (2021); Yang and
Karniadakis (2020). The MFG formulation of generative modeling in this paper
allows us to incorporate HJB regularizers to _any_ generative model that can
be written as a MFG, such as score-based generative models, in a systematic
fashion. Using HJB regularizers is a way to explicit encode the inductive
biases of generative flows during their training.

After reviewing HJB regularizers for normalizing flows in Section 6.2, we
formulate an HJB regularizer for score-based generative models in Section 6.3
and report the related numerical experiments in Section 6.4.

###  6.1 Hamilton-Jacobi-Bellman regularizers

Notice that the following penalty term can be added to the potential MFG
objective function, $\mathcal{J}(v,\rho)$ in (11), or any of its equivalent
forms, without changing the true minimizer,

$\displaystyle\mathcal{R}(U,\rho)=$
$\displaystyle\alpha_{1}\int_{0}^{T}\int_{\mathbb{R}^{d}}\left|\frac{\partial
U}{\partial t}-H(x,\nabla U)\right|\rho(x,t)\,\text{d}x\,\text{d}t$ (81)
$\displaystyle+\alpha_{2}\int_{\mathbb{R}^{d}}\left|U(x,T)-M(x,\rho(\cdot,T))\right|\rho(x,T)\,\text{d}x,$

where $\alpha_{1},\alpha_{2}>0$. Recall that $v$ is related to $U$ through
$v=-\nabla_{p}H(x,\nabla U)$, so we may also minimize the MFG objective in
(11) over $U$. Notice that the first term enforces the solution of the HJB
over $\mathbb{R}^{d}\times[0,T]$, while the second term enforces the terminal
condition. In essence, since we know any minimizer of the potential MFG
necessarily satisfies the HJB equation, the addition of the HJB regularizer
penalizes solutions that violate the optimality conditions and may help find
the minimizer faster.

###  6.2 HJB regularizers for normalizing flows

####  6.2.1 Optimal-transport normalizing flows

In addition to introducing the optimal transport cost to the normalizing flow,
Onken et al. (2021) empirically found that incorporating the HJB optimality
conditions as an additional regularizer to the optimization objective improved
training performance. They also found that the solutions to be more regular,
in the sense that particle trajectories were “straighter.” To be precise, the
optimality conditions (28) defines the regularization functional to be

$\displaystyle\mathcal{R}(U,\rho)=$
$\displaystyle\alpha_{1}\int_{0}^{T}\int_{\mathbb{R}^{d}}\left|\frac{\partial
U}{\partial t}-\frac{1}{2}|\nabla
U|^{2}\right|\rho(x,t)\,\text{d}x\,\text{d}t$ (82)
$\displaystyle+\alpha_{2}\int_{\mathbb{R}^{d}}\left|U(x,T)-\left(1+\log\frac{\rho(x,T)}{\rho_{ref}(x)}\right)\right|\rho(x,T)\,\text{d}x,$

for some $\alpha_{1},\alpha_{2}\geq 0$. This regularizer does not change the
true minimizer of the MFG (27) as they share the same minimizer. In Onken et
al. (2021), they do not enforce the terminal condition, i.e., $\alpha_{2}=0$,
as it requires the ability to evaluate the data distribution $\pi(x)$. This
regularizer can be cheaply evaluated since $U$ is exactly the potential
function that is learned, and so only derivatives of $U$ are needed to
evaluate the functional.

The idea of using the HJB as a regularizer has been introduced previously in
Ruthotto et al. (2020); Yang and Karniadakis (2020). In Ruthotto et al.
(2020), the authors solved a mean-field game with the tools of normalizing
flows, and chose to include HJB regularizer to improve stability when solving
the optimization problem. In Yang and Karniadakis (2020), GANs and flow-based
generative models are learned with the regularization functional. This
regularizer is a weak constraint since it is integrated over space and time
rather than being enforced pointwise. The inclusion of a partial differential
equation in a loss function in this weak sense is similar to the objective
function for training physics-informed neural networks (PINNs) Raissi et al.
(2019).

####  6.2.2 Normalizing flows with bounded velocities

We can use the results of Corollary 4 to formulate an HJB regularizer for the
canonical formulation of normalizing flows with bounded velocities. The
primary difference is the form of the HJB equation; the terminal condition is
the same as for the OT-flow in (82)

$\displaystyle\mathcal{R}(U,\rho)=$
$\displaystyle\alpha_{1}\int_{0}^{T}\int_{\mathbb{R}^{d}}\left|\frac{\partial
U}{\partial t}-c|\nabla U|\right|\rho(x,t)\,\text{d}x\,\text{d}t$ (83)
$\displaystyle+\alpha_{2}\int_{\mathbb{R}^{d}}\left|U(x,T)-\left(1+\log\frac{\rho(x,T)}{\rho_{ref}(x)}\right)\right|\rho(x,T)\,\text{d}x.$

###  6.3 HJB regularizers for score-based generative modeling and beyond

The use of an HJB regularizer during the training of a generative model is not
limited to the normalizing flows. Here we devise an HJB regularizer, discussed
in Section 6.1, for score-based generative models. Recall that for SGMs, the
score function $\mathsf{s}(y,s)=\nabla\log\eta(y,s)$ is learned via score-
matching. Furthermore, the score function is equal to the negative gradient of
solution to the Hamilton-Jacobi-Bellman equation (39) which comes from the SGM
MFG (37) optimality conditions. We cannot use the HJB equation as a
regularizer directly here since, in practice, each component of the score
function is learned individually and so the approximate score function will
not exactly be the gradient of a scalar function. However, by taking the
gradient of the HJB, it is possible to derive a regularizer for the score
function directly:

$\displaystyle\mathcal{R}_{1}(\mathsf{s})$
$\displaystyle=\alpha_{1}\int_{\mathbb{R}^{d}}\int_{0}^{T}\left|\frac{\partial\mathsf{s}}{\partial
s}-\nabla(f^{\top}\nabla\mathsf{s})-\frac{\sigma^{2}}{2}\nabla|\mathsf{s}|^{2}-\nabla(\nabla\cdot
f)-\frac{\sigma^{2}}{2}\nabla(\nabla\cdot\mathsf{s})\right|^{p}\eta(y,s)\,\text{d}s\,\text{d}y$
(84)
$\displaystyle+\alpha_{2}\int_{\mathbb{R}^{d}}|\mathsf{s}(y,0)-\nabla\log\pi(y)|^{2}\eta(y,0)\,\text{d}y,$

where $\eta(y,s)$ is given by (47) and $p\geq 1$. Recall that by
reparametrizing time, the terminal condition is instead an initial condition.
Moreover, observe that enforcing the terminal condition amounts to adding an
explicit score-matching objective for the data distribution Song and Ermon
(2019)! In particular, if the norm is chosen to be the Euclidean norm, then
the terminal condition can be enforced by adding an additional implicit score-
matching objective for the initial distribution $\eta(\cdot,0)=\pi$. In
particular, observe that

$\displaystyle\mathcal{R}_{1}(\mathsf{s})$
$\displaystyle=\alpha_{1}\sum_{i=1}^{d}\int_{\mathbb{R}^{d}}\int_{0}^{T}\left|\frac{\partial\mathsf{s}^{(i)}}{\partial
s}-\frac{\partial f^{\top}}{\partial
x_{i}}\mathsf{s}-f^{\top}\frac{\partial\mathsf{s}}{\partial
x_{i}}-\sigma^{2}\mathsf{s}^{\top}\frac{\partial\mathsf{s}}{\partial
x_{i}}-\nabla\cdot\frac{\partial f}{\partial
x_{i}}-\frac{\sigma^{2}}{2}\Delta\mathsf{s}^{(i)}\right|^{p}\eta(y,s)\,\text{d}s\,\text{d}y$
(85)
$\displaystyle+\alpha_{2}\int_{\mathbb{R}^{d}}\left(|\mathsf{s}(y,0)|^{2}+2\nabla\cdot\mathsf{s}(y,0)\right)\pi(y)\,\text{d}y,$

where $\mathsf{s}^{(i)}$ denotes the $i$-th component of the score function.
In Section 6.4, we will demonstrate the effect of this particular regularizer
on a simple example.

Alternatively, if we implement a version of SGM that learns the log-density
instead of the score function, we can use the HJB regularizer directly like in
the OT-Flow in Section 6.2. In this case, a scalar function $\varphi(y,s)$ is
learned using the score matching objective, where
$\mathsf{s}(y,s)=\nabla\varphi(y,s)$, and $U(y,s)=-\varphi(y,s)$. Moreover,
since the terminal constraint cannot be imposed directly, as we do not have
access to evaluations of $\pi(y)$, we can again use the score-matching
objective:

$\displaystyle\mathcal{R}_{2}(\varphi)$
$\displaystyle=\alpha_{1}\int_{\mathbb{R}^{d}}\int_{0}^{T}\left|\frac{\partial\varphi}{\partial
s}-f\cdot\nabla\varphi-\frac{\sigma^{2}}{2}|\nabla\varphi|^{2}-\nabla\cdot
f-\frac{\sigma^{2}}{2}\Delta\varphi\right|\eta(y,s)\,\text{d}s\,\text{d}y$
(86)
$\displaystyle+\alpha_{2}\int_{\mathbb{R}^{d}}\left(|\nabla\varphi(y,0)|^{2}+2\Delta\varphi(y,0)\right)\pi(y)\,\text{d}y.$

In contrast, for the OT-flow discussed in Section 6.2, the terminal constraint
cannot be applied as it requires evaluations of the target log-density.
Learning such a scalar function via score-matching has been studied previously
Song and Kingma (2021).

A similar PDE-based regularizer for SGM was studied in Lai et al. (2023) where
the uncontrolled Fokker-Planck equation is used to derive a PDE for the score
function. A crucial difference is that the regularizer presented in Lai et al.
(2023) does not enforce the terminal condition associated with the PDE
problem. In our paper, we can impose the terminal condition with little
additional computational cost by score-matching. In Section 6.4 we show that
the inclusion of the terminal cost can dramatically affect the quality of the
generative model.

######  Remark 17.

In this paper, we have focused on the training of score-based generative
models by solving a mean-field game. On the other hand, we may also study
mean-field games that can be solved through score-matching. Similar to how
Ruthotto et al. (2020); Huang et al. (2022) use tools for training normalizing
flows to solve mean-field games, we may also study how score-matching can
solve certain mean-field games and its associated high-dimensional PDEs. For
example, the use of entropy-based regularizers for solving mean-field games
have also been recently explored Guo et al. (2022). In particular, one could
solve certain high-dimensional HJB equation with only samples of the terminal
data using score-based generative modeling methods.

###  6.4 A numerical demonstration: HJB-regularized SGM

We demonstrate the effects of the HJB regularizer for score-based generative
modeling discussed in Section 6.4. We will highlight numerical and algorithmic
aspects one could explore in the future. The results here constitute only a
preliminary investigation of the potential computational contributions of MFG-
based generative modeling.

In this example, we choose $f(x,t)=\frac{x}{2}$ and $\sigma(t)=1$, so that the
induced noising process is an Ornstein-Uhlenbeck (OU) process. Recall from
Section 6.3 that the regularizer for this example is

$\displaystyle\mathcal{R}_{1}(\mathsf{s})=$
$\displaystyle\alpha_{1}\sum_{i=1}^{d}\int_{\mathbb{R}^{d}}\int_{0}^{T}\left|\frac{\partial\mathsf{s}^{(i)}}{\partial
s}-\frac{\partial f^{\top}}{\partial
x_{i}}\mathsf{s}-f^{\top}\frac{\partial\mathsf{s}}{\partial
x_{i}}-\mathsf{s}^{\top}\frac{\partial\mathsf{s}}{\partial
x_{i}}-\frac{1}{2}\Delta\mathsf{s}^{(i)}\right|^{p}\eta(y,s)\,\text{d}s\,\text{d}y$
(87)
$\displaystyle+\alpha_{2}\int_{\mathbb{R}^{d}}\left(|\mathsf{s}(y,0)|^{2}+2\nabla\cdot\mathsf{s}(y,0)\right)\pi(y)\,\text{d}y,$

where for some $p\geq 1$. In this example, we choose $p=1$ or $p=2$. The
regularized implicit score-matching problem is then

$\displaystyle\min_{\theta}\alpha_{0}\int_{0}^{T}\int_{\mathbb{R}^{d}}\left(\frac{1}{2}|\mathsf{s}_{\theta}(y,s)|^{2}+\nabla\cdot\mathsf{s}_{\theta}(y,s)\right)\eta(y,s)\,\text{d}y\,\text{d}s+\mathcal{R}_{1}(\mathsf{s}_{\theta}).$
(88)

The data distribution $\pi(y)$ is the checkerboard distribution, see Figure
2(a). The checkerboard distribution presents particular challenges to score-
based generative models, mainly in that the distribution has disconnected
support and the density is discontinuous.

We informally investigate and observe the impact of different choices of the
HJB regularization parameters, $\alpha_{0}$, $\alpha_{1}$, $\alpha_{2}$, and
$p$. For all examples, the total simulation time is $T=3$, and the reverse SDE
is simulated with an Euler-Maruyama scheme with $\Delta t=0.001.$ The forward
SDE is an OU process and can be simulated exactly. For all experiments, the
score function
$\mathsf{s}_{\theta}:\mathbb{R}^{2}\times\mathbb{R}\to\mathbb{R}^{2}$ is a
neural network with two hidden layers, with 32 nodes per hidden layer. We use
the Gaussian error linear unit (GeLU) activation function, and train the
network over $10^{5}$ batches of 64 samples, with learning rate $0.001$. The
derivatives are computed using the DeepXDE package Lu et al. (2021).

(a) True samples
filler space filler space

(b) Without any HJB regularization

(c) HJB eq., no terminal. $p=1$

(d) HJB eq., no terminal. $p=2$

(e) Terminal constraint only

(f) HJB eq. $+$ terminal, $p=1$

(g) HJB eq. $+$ terminal, $p=2$

(h) HJB eq. $+$ terminal, $p=2$, smaller $\alpha_{i}$

Figure 2: Effects of different choices of HJB regularization on score-based
generative modeling. In these examples, $\alpha_{0}=1$ in (88).

We can clearly see that the HJB regularizer has a nontrivial effect on the
generative properties of the score-based generative modeling. Cases where the
terminal constraint is enforced such as Figures 2(e) and 2(f), yields sharper
boundaries. However, they also seem to exhibit Gibbs phenomenon Strauss
(2007), in which more samples lie on boundaries of the checkerboard density.
In Figure 2(h), we see that reducing the constant $\alpha_{2}$ for the
terminal constraint can help make the boundaries sharper, without exhibiting
as much Gibbs phenomenon. Comparing Figures 2(f) and 2(g), we find different
behavior when choosing different values of $p$. Figure 2(h) shows that the
parameters can be tuned to produce an improved generative model. On the other
hand, if only the HJB equation is enforced as presented in Lai et al. (2023),
then the regularizers may produce worse results than if no regularizer was
present. This is likely due the fact that without the terminal constraint, the
function $\mathcal{s}(x,t)=0$ minimizes the regularizer. Therefore, the use of
the HJB regularizer alone skews the learned score function towards zero.

These simple examples clearly show that HJB regularization has the potential
to provide improvements to score-based methods. They also demonstrate that a
more extensive numerical and algorithmic investigation on the effects of HJB
regularization strategies in MFG-based generative modeling are warranted for
future research.

######  Remark 18.

We note that we are not restricted to use the distribution $\eta(y,s)$ when
evaluating the HJB regularizer (87). Using $\eta(y,s)$ as the measure the HJB
equations are integrated with respect to is a natural choice since it is
already used in the implicit score-matching objective. The choice of measure
here is, however, arbitrary; for example, Lu et al. (2021) chooses a uniform
distribution over space and time. The effects of the choice of measure can
also be another direction for future research.

###  6.5 Score-matching via PINNs

The HJB regularizer provides an alternate way for implementing score-based
generative modeling. Since the terminal conditions (46) can be imposed
directly, we can recover results similar to standard score-based generative
modeling by _only_ minimizing the HJB regularizer, i.e., $\alpha_{0}=0$ in
(88). In Section 4, we established that score-based generative modeling is
based on solving a pair of partial differential equations. We can therefore
obtain approximate solutions of the optimality conditions by minimizing the
regularizers directly, which we can then use with the reverse SDE (34) to
generate new samples. Minimizing the HJB regularizer directly is the same task
as training a physics-informed neural network (PINN) Raissi et al. (2019); Lu
et al. (2021).

In Figure 3 we show results of score-based generative modeling by only
minimizer the HJB regularizer. We see different performance depending on the
choice of parameter values. Comparing Figure 3(a) with 3(b), we see that when
the terminal constraint has a larger constant than the HJB equation, the
generated distribution has sharper edges. When $p=1$, we see many particles
end up on the edges, which is indicative of a Gibbs phenomenon. We hope the
results in this section may spawn further research regarding the use of HJB
regularizers for generative modeling.

(a) HJB regularizer only with
$p=2,\,\alpha_{1}=1$

(b) HJB regularizer only with
$p=2,\,\alpha_{1}=0.1$

(c) HJB regularizer only with
$p=1$

Figure 3: Score-based generative modeling via PINNs. Here $\alpha_{0}=0$ in
(88). Compare these figures with standard SGM in Figure 2(a).

##  7 Inventing generative models using MFGs

We provide some examples of how the MFG framework can easily lead to new
variants of generative models, as well as enhancements of existing ones in a
coherent and motivated fashion. The richness of potential generative modeling
formulations via the MFG framework is partly captured in Table 3 but it is
primarily a topic for future research. The examples we discuss here are by no
means exhaustive. We once again emphasize that while one can conjure the
following generative models with Table 3, further numerical and algorithmic
investigation is needed to evaluate their practical utility.

###  7.1 Optimal transport BG: a corrected Boltzmann generator

As we discussed in Section 3.3.2, Boltzmann generators build generative models
using normalizing flows that help when sampling from probability distributions
that arise in computational chemistry Noé et al. (2019). The normalizing flow
formulation in (31) is, however, ill-posed without additional constraints or
structure as described in Section 3.2. One (of many) possible fixes is to
introduce a quadratic running cost like in the OT-flow in Section 3.3.1. We
introduce the optimal transport Boltzmann generator MFG,

$\displaystyle\inf_{v,\rho}\left\\{\lambda\mathcal{D}_{KL}(\pi\|\rho(\cdot,T))+(1-\lambda)\mathcal{D}_{KL}(\rho(\cdot,T)\|\pi)+\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)\,\text{d}x\,\text{d}t\right\\}$
(89) $\displaystyle\frac{\partial\rho}{\partial
t}+\nabla\cdot(v\rho)=0,\,\rho(x,0)=\rho_{0}(x)$

for some $\lambda\in[0,1]$. In this setting, the Hamiltonian will be
quadratic, so the optimal velocity field will be the gradient of a scalar
function. Therefore, like in the OT-flow, we may optimize for a single scalar
function, rather than for each component of a vector-valued velocity function.
Moreover, in Köhler et al. (2020), the normalizing flow velocity field was
parametrized as the gradient of a scalar function, even though the
optimization problem did not include a quadratic running cost. Adding such a
running cost will mathematically justify the approach taken in Köhler et al.
(2020), and as a result, likely facilitate the algorithm to learn a
normalizing flow with _gradient_ structure. Furthermore, like in Section 6.3,
we may introduce an HJB regularizer to enhance training of the Boltzmann
generator.

###  7.2 Stochastic normalizing flows via MFGs

Since the introduction of continuous normalizing flows based on deterministic
dynamics, there have been several attempts to extend it to stochastic
differential equations Wu et al. (2020); Hodgkinson et al. (2020); Zhang and
Chen (2021). While each devises their own computational approaches to and
implementation of a stochastic normalizing flow, it is unclear if their
mathematical formulations are well-posed. With the MFG formulation, it is
simple to devise well-posed stochastic normalizing flows. For example,
consider the MFG

$\displaystyle\inf_{v,\rho}\left\\{\mathcal{D}_{KL}(\pi\|\rho(\cdot,T))+\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)\,\text{d}x\,\text{d}t\right\\}$
(90) $\displaystyle\frac{\partial\rho}{\partial
t}+\nabla\cdot(v\rho)=\frac{\sigma^{2}}{2}\Delta\rho,\,\rho(x,0)=\rho_{0}(x).$

We have already studied the well-posedness of this mean-field game: it has the
same optimality conditions as score-based generative models (39), except that
the terminal condition is based on the KL divergence instead of cross-entropy,
i.e.,

$\displaystyle U(x,T)=-\frac{\pi(x)}{\rho(x,T)}.$ (91)

One can then devise computational schemes for solving this MFG and be assured
that there exist consistent discretizations since the optimization problem is
well-defined. To further belabor the point about the HJB regularizer, we may,
again consider using the HJB regularizer for this newly introduced generative
model as well.

Moreover, as we have discussed in Remark 12, score-based generative models can
be considered to be a stochastic normalizing flow. In particular, if the
terminal cost $\mathcal{D}_{KL}(\pi\|\rho)$ is replaced with the cross-entropy
$-\mathbb{E}_{\rho}[\log\pi]$, (90) is identical to the SGM MFG (37) for
$f\coloneqq 0$, and $\rho_{0}(x)=\eta(x,T)$ defined by (41).

###  7.3 SGM and the cross-entropy regularization of optimal transport

The Schrödinger bridge problem in Section 4.5 is often studied as an
entropically regularized version of the optimal transport problem Chen et al.
(2021). In fact, if we replace the stochastic dynamics in (63) by the Fokker-
Planck equation, we obtain a potential formulation (11) of the SBP. Let
$\Omega\subset\mathbb{R}^{d}$ be compact, and consider the MFG

$\displaystyle\inf_{v,\rho}\left\\{\int_{0}^{T}\int_{\Omega}|v(x,t)|^{2}\rho(x,t)\,\text{d}x\,\text{d}t:\frac{\partial\rho}{\partial
t}+\nabla\cdot(v\rho)=\frac{\sigma^{2}}{2}\Delta\rho,\,\rho(x,0)=\rho_{0}(x),\,\rho(x,T)=\pi(x)\right\\}.$
(92)

When $\sigma=0$, this MFG is the Benamou-Brenier formula, which is the dynamic
formulation of the optimal transport problem Villani et al. (2009);
Santambrogio (2015). This is the problem of transporting a distribution
$\rho_{0}$ to $\pi$ with minimal cost Santambrogio (2015). When $\sigma>0$, it
is a regularization of the optimal transport problem Santambrogio (2015); Chen
et al. (2021).

Since the terminal constraint is difficult to enforce pointwise, it is common
to relax it by introducing a penalty term Genevay et al. (2017). If we replace
the terminal constraint with a cross-entropy loss function
$-\mathbb{E}_{\rho}[\log\pi]$, we can interpret the resulting MFG as a _cross-
entropy regularized_ entropic dynamic optimal transport problem:

$\displaystyle\inf_{v,\rho}\left\\{-\int_{\Omega}\rho(x,T)\log\pi(x)\,\text{d}x+\int_{0}^{T}\int_{\Omega}|v(x,t)|^{2}\,\text{d}x\,\text{d}t:\frac{\partial\rho}{\partial
t}+\nabla\cdot(v\rho)=\frac{\sigma^{2}}{2}\Delta\rho,\,\rho(x,0)=\rho_{0}(x)\right\\}.$
(93)

Notice, however, this is exactly equivalent to the SGM MFG (37) in Theorem 7
for $f\coloneqq 0$. With this observation, we may approximate solutions to a
certain class of optimal transport problems with computational tools of score-
based generative modeling, which have proved to be extremely efficient in
high-dimensions.

###  7.4 Interpolating between MFGs: from Wasserstein gradient flows to
normalizing flows and Wasserstein geodesics

We refer back to the connection between Wasserstein gradient flows and mean-
field games described by Theorem 14. In Section 5, we studied the MFG (69) in
the asymptotic regime where $\epsilon\to 0$. Notice, however, there is another
interesting regime when $\epsilon\to\infty$, in which the MFG (69) converges
to a regularized optimal transport problem:

$\displaystyle\min_{v,\rho}\left\\{\mathcal{F}(\rho(\cdot,T))+\int_{0}^{T}\int_{\Omega}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)\,\text{d}x\,\text{d}t\right\\}\quad
s.t.\quad\frac{\partial\rho}{\partial
t}+\nabla\cdot(v\rho)=0,\,\rho(x,0)=\rho_{0}(x).$ (94)

Since the functional $\mathcal{F}(\rho)$ is typically some divergence from
$\rho$ to target $\pi$, for example,
$\mathcal{F}(\rho)=\mathcal{D}_{KL}(\rho\|\pi)$, this problem is a
generalization of the OT normalizing flow in Section 3.3.1. Recall, however,
that (94) is the limiting case of (69) where $\epsilon\to\infty$, which
implies that for $\epsilon\in(0,\infty)$ we have MFGs that _interpolate_ the
Wasserstein gradient flow, and the OT-flow. Following the same steps as in
Section 3.3.1, the optimality conditions of the MFG (94) are

$\displaystyle\begin{dcases}&-\frac{\partial U}{\partial t}+\frac{1}{2}|\nabla
U|^{2}=0\\\ &\frac{\partial\rho}{\partial t}-\nabla\cdot\left(\rho\nabla
U\right)=0\\\
&U(x,T)=\frac{\delta\mathcal{F}}{\delta\rho}(\rho(\cdot,T)),\,\rho(x,0)=\rho_{0}(x).\end{dcases}$
(95)

Moreover, recall from Section 7.3 that $\mathcal{F}(\rho)$ in the terminal
cost can be interpreted as a relaxation of a terminal constraint
$\rho(x,T)=\pi(x)$. Therefore, (94) is a relaxation of the Benamou-Brenier
dynamic formulation of optimal transport. Since solutions of the optimal
transport problem are geodesics connecting $\rho_{0}$ and $\pi$ in Wasserstein
space Villani et al. (2009); Santambrogio (2015), solutions of the MFG (69)
for $\epsilon\in(0,\infty)$ can be interpreted as interpolating between
Wasserstein gradient flows ($\epsilon\to 0$) and Wasserstein geodesics
($\epsilon\to\infty$).

##  8 A modular mean-field games framework for generative modeling

|  Mean-field game (11) |  |  |  |
---|---|---|---|---|---
Model | $\mathcal{M}(\rho)$ | $\mathcal{I}(\rho)$ | $L(x,v)$ | Dynamics | $H(x,p)$ |  Optimal $v^{*}$ |  HJB Reg? | Refs
Continuous normalizing flow | $\mathcal{D}_{KL}(\rho\|\rho_{ref})$ | 0 | 0 | $\,\text{d}x=v\,\text{d}t$ | $\sup_{v\in K}-p^{\top}v$ | $-\nabla_{p}H(x,\nabla U)$ | $\Large{\usym{2717}}$ | 1
Score-based generative modeling | $-\mathbb{E}_{\rho}\left[\log\pi\right]$ | 0 | $\frac{|v|^{2}}{2}-\nabla\cdot f$ |  $\,\text{d}x=(f+\sigma v)\,\text{d}t$ $+\sigma\,\text{d}W_{t}$ |  $-f\cdot p+\frac{\sigma^{2}}{2}|p|^{2}$ $+\nabla\cdot f$ | $\sigma\nabla\log\eta_{T-t}$ | $\Large{\usym{2717}}$ | 2
Score-based probability flow | $-\frac{1}{2}\mathbb{E}_{\rho}\left[\log\pi\right]$ | $\mathbb{E}_{\rho}\left[\frac{|\sigma\nabla\log\rho|^{2}}{8}\right]$ | $\frac{|v|^{2}}{2}-\frac{\nabla\cdot f}{2}$ | $dx=(f+\sigma v)\,dt$ |  $-f\cdot p+\frac{\sigma^{2}}{2}|p|^{2}$ $+\nabla\cdot f$ | $\frac{\sigma}{2}\nabla\log\eta_{T-t}$ | $\Large{\usym{2717}}$ | 3
Wasserstein gradient flow (WGF) $(\epsilon\to 0)$ | $\mathcal{F}(\rho)e^{-T/\epsilon}$ | $\frac{e^{-t/\epsilon}}{\epsilon}\mathcal{F}(\rho)$ | $\frac{e^{-t/\epsilon}}{2}|v|^{2}$ | $\,\text{d}x=v\,\text{d}t$ | $\frac{1}{2}e^{t/\epsilon}|p|^{2}$ | $-\nabla\frac{\delta\mathcal{F}}{\delta\rho}$ | $\Large{\usym{2717}}$ | 4
OT-Flow | $\mathcal{D}_{KL}(\rho\|\rho_{ref})$ | 0 | $\frac{1}{2}|v|^{2}$ | $\,\text{d}x=v\,\text{d}t$ | $\frac{1}{2}|p|^{2}$ | $-\nabla U$ | $\Large{\usym{2713}}$ | 5
Boltzmann generator |  $\lambda\mathcal{D}_{KL}(\pi\|\rho)$ $+(1-\lambda)\mathcal{D}_{KL}(\rho\|\pi)$ | 0 | 0 | $\,\text{d}x=v\,\text{d}t$ | $\sup_{v\in K}-p^{\top}v$ | $-\nabla_{p}H(x,\nabla U)$ | $\Large{\usym{2717}}$ | 6
Schrödinger bridge | $\rho=\pi$ | 0 | $\frac{1}{2}|v|^{2}$ | $\,\text{d}x=\sigma v\,\text{d}t+\sigma\,\text{d}W_{t}$ | $\frac{1}{2}|p|^{2}$ | $-\sigma\nabla U$ | $\Large{\usym{2717}}$ | 7
Generalized Schrödinger bridge | $\rho=\pi$ | $\mathcal{I}(x,\rho)$ | $\frac{1}{2}|v|^{2}$ |  $\,\text{d}x=\sigma v\,\text{d}t+\sigma\,\text{d}W_{t}$ | $\frac{1}{2}|p|^{2}$ | $-\sigma\nabla U$ | $\Large{\usym{2717}}$ | 8
HJB-regularized SGM | $-\mathbb{E}_{\rho}[\log\pi]$ | 0 | $\frac{|v|^{2}}{2}-\nabla\cdot f$ |  $\,\text{d}x=(f+\sigma v)\,\text{d}t$ $+\sigma\,\text{d}W_{t}$ |  $-f\cdot p+\frac{\sigma^{2}}{2}|p|^{2}$ $+\nabla\cdot f$ | $\sigma\nabla\log\eta_{T-t}$ | $\Large{\usym{2713}}$ | 9
Stochastic OT normalizing flow | $\mathcal{D}_{KL}(\pi\|\rho)$ | 0 | $\frac{1}{2}|v|^{2}$ |  $\,\text{d}x=\sigma v\,\text{d}t+\sigma\,\text{d}W_{t}$ | $\frac{1}{2}|p|^{2}$ | $-\sigma\nabla U$ | $\Large{\left(\usym{2713}\right)}$ | 10
OT-Boltzmann generator |  $\lambda\mathcal{D}_{KL}(\pi\|\rho)$ $+(1-\lambda)\mathcal{D}_{KL}(\rho\|\pi)$ | 0 | $\frac{1}{2}|v|^{2}$ | $\,\text{d}x=v\,\text{d}t$ | $\frac{1}{2}|p|^{2}$ | $-\nabla U$ | $\Large{\left(\usym{2713}\right)}$ | 11
Generalized OT-Flow (Relaxed WGF $\epsilon>0$) | $\mathcal{F}(\rho)e^{-T/\epsilon}$ | $\frac{e^{-t/\epsilon}}{\epsilon}\mathcal{F}(\rho)$ | $\frac{e^{-t/\epsilon}}{2}|v|^{2}$ | $\,\text{d}x=v\,\text{d}t$ | $\frac{1}{2}e^{t/\epsilon}|p|^{2}$ | $-\nabla U$ | $\Large{\left(\usym{2713}\right)}$ | 12
Build your own generative model |  Choose your |  own cost |  functions and |  dynamics here | ? | ? | ? | ?

Table 3: Summary of generative modeling and related topics as mean-field
games.
1\. Chen et al. (2018); Grathwohl et al. (2018) & Sec. 3.2
2\. Song et al. (2021b); De Bortoli et al. (2021) & Sec. 4
3\. Maoutsa et al. (2020); Song et al. (2021b) Sec. 4.4
4\. Santambrogio (2015); Rossi et al. (2011); Arbel et al. (2019); Gu et al.
(2022) & Sec. 5.1
5\. Onken et al. (2021); Huang et al. (2022) & Sec. 3.3.1
6\. Noé et al. (2019); Köhler et al. (2020) & Sec. 3.3.2
7\. Pavon et al. (2021); Chen et al. (2021) & Sec. 4.5
8\. Liu et al. (2022) & Sec. 4.5
9\. New & Sec. 6.3, Sec. 6.4, 10. New & Sec. 7.2
11\. New & Sec. 7.1, 12. New & Sec. 7.4

##  9 Discussion and outlook

In this paper, we showed that mean-field games provide a natural framework for
studying continuous-time generative models. Given samples from a target
distribution, continuous-time generative models aim to approximate the target
by building a dynamical system that evolves samples from a reference
distribution to those of the target. Depending on the choice of the
continuous-time model, the dynamics of each particle may depend on the
evolution of other particles or on the final state of the density of
particles. The mean-field games framework can be readily used as an
organizational principle for understanding the generative modeling task.

We showed how normalizing flows, score-based generative models, and
Wasserstein gradient flows can be derived from different MFGs. Moreover, we
demonstrated how MFGs can explain and enhance generative models by revealing
and exploiting their mathematical structure, yielding new computational tools,
and provide a system for inventing new generative models. We summarize our
main contributions here:

  * •

Explaining and categorizing generative models: We showed that the potential
formulation of MFG in Section 2.2 describes major classes of generative
models. We showed how with different choices cost function, mean-field games
correspond to different generative models. Each generative model can be
defined by four ingredients: the terminal, interaction, and running cost
functions, and the form of the dynamical system. The MFG optimality conditions
uncover the common mathematical structure of all flow and diffusion-based
generative models. While the Fokker-Planck equation is clearly the generative
portion of the model, the Hamilton-Jacobi-Bellman equation can be interpreted
as the ‘discriminative’ part, which determines the optimal velocity field of
the generative model. The optimality conditions describe the mathematical
structure that identifies the inductive bias of generative flows. Overall, the
MFG perspective yields a simple way of organizing the vast collection of flow
and diffusion-based generative models. We exhibit some of these methods in
Table 3. The table also provides a way to understand generative models in
relation to each other.

Furthermore, the MFG formulation provides insight into the _well-posedness_ of
a generative model. In particular, well-posedness of a generative model can be
studied in terms of the well-posedness of the optimality conditions, which are
a set of partial differential equations. In Section 3.2 we used the optimality
conditions of MFGs to show that the canonical formulation of normalizing flows
is ill-posed because the Hamiltonian, and therefore the associated HJB
equation, is not well-defined without additional constraints to the problem.
The MFG formulation then provides various options to make the problem well-
posed. In Section 5.1, we showed that the finite-time relaxations of the
Wasserstein gradient flow require the terminal and interaction costs to be
related so that the flow avoids boundary layers at terminal time.

  * •

Enhancing generative models: Studying generative models within the MFG
framework provides access to optimality conditions which describe the solution
to the generative modeling problem. These HJB regularizers, while first
developed for optimal transport flows, can be applied to any generative model
described by the MFG framework presented here. In Sections 6.3 and 6.4 we
demonstrated how these HJB regularizers may be applied to improve score-based
generative modeling. The HJB regularizers provides one approach for enforcing
the inductive bias of generative models, which may accelerate their training.

Moreover, the MFG perspective provides ideas for better relaxations of
problems in optimal transport. For example, in Section 7.3 we remarked that
the entropy-regularized Benamou-Brenier dynamic formulation of optimal
transport, where the terminal constraint is substituted with a terminal cross-
entropy cost function, is equivalent to score-based generative modeling. Thus,
the MFG perspective provides insight into how to choose relaxations that are
more easily solved via machine learning techniques.

  * •

Inventing generative models: The MFG framework provides a systematic way for
inventing new generative models. In Section 7 we proposed three new generative
models that either extend existing approaches (Sections 6.3 and 7.2) or amend
ill-posed approaches (Section 7.1).

Our contributions here have only discussed a fraction of the potential
insights of each of the above three categories. The main goal of this paper
has been to illustrate the richness of studying generative models through the
mean-field game perspective. For future work, we provide a decidedly non-
exhaustive list of possible future research directions. Table 3, however,
beckons for further future investigations and contributions to generative
modeling and mean field games, well-beyond what is discussed in this paper.

  * •

The role of stochasticity: We can study further the well-posedness of
generative models by using the theory of MFGs and partial differential
equations. These aspects are important for machine learning applications as
they may lead to consistent discretizations of generative models, and more
stable algorithms for training such models. MFG optimality conditions may
provide insight into the role of stochasticity in generative models. For
stochastic dynamics (e.g., diffusion-based methods), gradients of the HJB are
always smooth. Therefore learning the optimal velocity is always well-defined.
This is not the case for deterministic dynamics such as normalizing flows
where gradients of the corresponding first-order HJB and hence optimal
velocities may have jumps, creating ambiguities in the optimal velocities that
the generative algorithm will need to somehow resolve. In this sense,
stochastic generative algorithms have theoretical and possibly practical
advantages, that need to be explored further using PDE and MFG regularity
theory.

  * •

Finite sample analysis: Mean-field games are asymptotic dynamics of a large
population of interacting agents. Studying MFGs in the nonasymptotic regime