# Combining Wasserstein-1 and Wasserstein-2 proximals: robust manifold
learning via well-posed generative flows

Hyemin Gu Markos A. Katouslakis Luc Rey-Bellet Benjamin J. Zhang
Department of Mathematics and Statistics
University of Massachusetts Amherst
<EMAIL_ADDRESS>
###### Abstract

We formulate well-posed continuous-time generative flows for learning
distributions that are supported on low-dimensional manifolds through
Wasserstein proximal regularizations of $f$-divergences. Wasserstein-1
proximal operators regularize $f$-divergences so that singular distributions
can be compared. Meanwhile, Wasserstein-2 proximal operators regularize the
paths of the generative flows by adding an optimal transport cost, i.e., a
kinetic energy penalization. Via mean-field game theory, we show that the
_combination_ of the two proximals is critical for formulating well-posed
generative flows. Generative flows can be analyzed through optimality
conditions of a mean-field game (MFG), a system of a backward Hamilton-Jacobi
(HJ) and a forward continuity partial differential equations (PDEs) whose
solution characterizes the optimal generative flow. For learning distributions
that are supported on low-dimensional manifolds, the MFG theory shows that the
Wasserstein-1 proximal, which addresses the HJ terminal condition, and the
Wasserstein-2 proximal, which addresses the HJ dynamics, are both necessary
for the corresponding backward-forward PDE system to be well-defined and have
a unique solution with provably _linear_ flow trajectories. This implies that
the corresponding generative flow is also unique and can therefore be learned
in a robust manner even for learning high-dimensional distributions supported
on low-dimensional manifolds. The generative flows are learned through
_adversarial training_ of continuous-time flows, which bypasses the need for
reverse simulation. We demonstrate the efficacy of our approach for generating
high-dimensional images without the need to resort to autoencoders or
specialized architectures.

##  1 Introduction

Continuous normalizing flows (CNF) are a class of generative models [1, 2]
based on learning deterministic ODE dynamics that transport a simple reference
distribution to a target distribution. However, two challenges arise when
applying CNFs for learning high-dimensional distributions: (a) the flows are
heavily discretization or implementation dependent due to lack of a unique
optimizer, and (b) they are not suitable for learning distributions supported
on manifolds, as is typically the case in many high-dimensional datasets, as
they need to be invertible in practice. For example, [3, 4] have noted that
without additional regularizers, normalizing flows can produce multiple valid
flows, many of which are difficult to integrate. [5] has noted that the test
accuracy of neural ODEs are discretization dependent. Moreover, as normalizing
flows need to be invertible, latent representations or specialized
architectures are frequently needed to successfully generate from
distributions that are supported on manifolds for even moderately high-
dimensional datasets [1, 2, 3, 4, 6]. The challenges of CNFs are fundamentally
rooted in the mathematical fact that their associated optimization problems
are _ill-defined_ , especially for distributions supported on low-dimensional
manifolds.

In this paper, through rigorous theory and numerical experiments, we
demonstrate that the composition of the Wasserstein-1 and Wasserstein-2
proximal regularizations is essential for formulating _well-posed_
optimization problems in the training of deterministic flow-based generative
models. Recent work [3, 6, 4] has shown that training generative flows via a
Kullback-Leibler (KL) divergence regularized with a Wasserstein-2 proximal
($\mathcal{W}_{2}$) through the Benamou-Brenier formulation of optimal
transport yields sample paths that are more regular. In fact, we prove via the
theory of Hamilton-Jacobi equations and mean-field games that regularizing the
KL divergence with a $\mathcal{W}_{2}$ proximal produces optimal trajectories
with _constant velocity_ , i.e., straight lines. The proof is based on
connections of Hamilton-Jacobi equations to the compressible Euler equations
in fluid mechanics, which lead to further interesting interpretations of
generative flows. Such flows can be discretized with fewer time steps, saving
on computational overhead and may learn the target distribution faster.

Learning distributions supported on manifolds, however, is still a challenge
with the Wasserstein-2 proximal alone as the KL (or, more generally, $f$-)
divergence is unable to compare distributions that are not absolutely
continuous with respect to each other as they are finite only for
distributions that share the same support. We propose regularizing the KL (or
$f$-) divergence with a Wasserstein-1 ($\mathcal{W}_{1}$) proximal as well [7,
8]. Applying a $\mathcal{W}_{1}$ proximal regularization of the KL allows us
to compare _mutually singular distributions_. The variational derivative of a
$\mathcal{W}_{1}$ regularized KL divergence is well-defined for _any_
perturbation in the space of probability distributions, including mutually
singular ones, meaning that the resulting training procedure (i.e., gradient
descent) is _robust_ for learning distributions supported on manifolds.

This composition of $\mathcal{W}_{1}$ and $\mathcal{W}_{2}$ proximals for
improving normalizing flows is grounded in the theory of mean-field games.
Generative flows can be roughly viewed as optimal control problems where a
velocity field, i.e., a control, is learned so that a collection of particles
will evolve from a reference distribution to a target distribution by
optimizing an objective functional involving interactions of generated and
real samples through various mean-field mechanisms that depend on the class of
models, e.g. CNFs, score-based models or gradient flows [9]. The _optimality
conditions_ of these MFGs are a pair of coupled backward-forward partial
differential equations (PDEs): a backward Hamilton-Jacobi equation whose
solution describes the optimal velocity field, and a forward continuity
equation describing the generation process. The backward-forward structure
mirrors that of forward and inverse flows in generative modeling [9]. The
terminal condition is associated with the target distribution and choice of
probability divergence. We propose a training objective with the composition
of $\mathcal{W}_{1}$ and $\mathcal{W}_{2}$ proximals whose well-posedness is
guaranteed by the well-posedness of the MFG optimality conditions. The
$\mathcal{W}_{2}$ proximal yields a well-defined Hamilton-Jacobi equation,
while the $\mathcal{W}_{1}$ proximal of $f$-divergences provides a well-
defined terminal condition for distributions supported on manifolds. We show
that the MFG is well-posed via a uniqueness theorem where we show that if the
solution of the MFG PDE system is smooth, then the solution is unique. This
uniqueness and well-posedness property of the generative flow produces more
_robust_ learning algorithms that are less dependent on implementation
details.

Figure 1: Optimality indicator (44) blows up without $\mathcal{W}_{1}$
proximal (left). The composition of $\mathcal{W}_{1}$ and $\mathcal{W}_{2}$
proximals minimizes kinetic energy and keeps the value less oscillatory while
training (right). See Section 6 for details.

Our algorithm is based on adversarial training, where the proximal structure
provides clear interpretations for the discriminators and generators. In
contrast to previous CNF approaches [3, 4, 2], the adversarial formulation
_avoids_ needing to invert the generative flow during training. Our resulting
$\boldsymbol{\mathcal{W}_{1}\oplus\mathcal{W}_{2}}$ proximal generative flows
robustly learns distributions supported on manifolds _without specialized
architectures or autoencoders_. We preview our numerical results in Figure 1,
which show the impact of $\mathcal{W}_{1}$ and $\mathcal{W}_{2}$ proximals for
training generative flows of distributions supported on manifolds. Notice that
without a $\mathcal{W}_{1}$ proximal, the terminal KL cost diverges over the
course of the training, while without a $\mathcal{W}_{2}$ proximal, the
kinetic energy oscillates during the course of the training, implying a lack
of convergence to a unique flow.

###  1.1 Contributions

  * •

We rigorously demonstrate that using the composition of Wasserstein-1 and
Wasserstein-2 proximals of $f$-divergences to train generative flows is
indispensable for _robustly_ learning distributions supported on low-
dimensional manifolds. The $\mathcal{W}_{2}$ proximal guarantees trajectories
are linear and the $\mathcal{W}_{1}$ proximal robustly learns distributions on
manifolds. These properties are due to the well-posedness of the optimization
problem, which is justified via the theory of mean-field games and the
properties of $\mathcal{W}_{1}$ and $\mathcal{W}_{2}$ proximals.

  * •

We produce an _adversarial training algorithm_ for training normalizing flows
based on the dual formulation of the Wasserstein-1 proximal of
$f$-divergences. This eliminates the need for forward-backward simulations of
trajectories that are needed to train other CNFs. The call our resulting
generative model $\boldsymbol{\mathcal{W}_{1}\oplus\mathcal{W}_{2}}$ proximal
generative flows.

  * •

We find that $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative flows
_robustly_ learn high-dimensional data supported on low-dimensional manifolds
_without_ pretraining autoencoders or tuning specialized architectures.

###  1.2 Related work

The use of Wasserstein proximals in generative modeling has already frequently
appeared in previous studies. Using $\mathcal{W}_{2}$ regularized KL
divergences for training CNFs were first introduced in [3] and [4]. The
quality of their image generation models, however, were still implementation
dependent. In [6], the flows are regularized with the Hamilton-Jacobi equation
corresponding to a $\mathcal{W}_{2}$ proximal, although their objective
function does not include the kinetic energy penalization. In all these
previous works, the HJ terminal conditions are ill-defined when the target
distributions are supported on low-dimensional manifolds, and therefore depend
on using some sort of autoencoder. $\mathcal{W}_{2}$ proximals of linear and
nonlinear functions were generally studied in [10]. Score-based generative
models were shown to be Wasserstein-2 proximals of cross-entropy in [11] using
ideas from [10] and [9]. Using adversarial training for normalizing flows was
considered in [12], but without a kinetic energy penalization (i.e., a
Wasserstein-2 proximal) a Jacobian normalization is employed to regularize
training. The key connection between mean-field games and generative flows [9]
is a crucial tool that is used to prove our proposed
$\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ are trained via well-posed optimization
problems.

The construction of Wasserstein gradient flow based the $\mathcal{W}_{1}$
proximal of $f$\- divergences for generative modeling were considered in [8].
There, a stable particle transport algorithm was constructed that is able to
flow between two mutually singular empirical distributions. Additional results
about Wasserstein-1 regularized $f$\- and Rényi divergences can be found in
[7, 13]. Their use for designing generative adversarial nets was studied in
[14].

##  2 Wasserstein proximals

Let $\mathcal{P}_{p}(\mathbb{R}^{d})$ be the Wasserstein-$p$ space, the space
of probability distributions on $\mathbb{R}^{d}$ with finite $p$-th moments
with the Wasserstein-$p$ distance. Wasserstein-1 and Wasserstein-2 proximals
have been been used intentionally or inadvertently for studying and designing
generative models [7, 13, 3, 8, 10, 4, 6, 11]. Wasserstein-$p$ proximals for
$p=1,2$ regularize cost functions (e.g., divergences)
$\mathcal{F}:\mathcal{P}_{p}(\mathbb{R}^{d})\rightarrow\mathbb{R}$ by an
infimal convolution between the function $\mathcal{F}(R)$ and the $p$-th power
of Wasserstein-$p$ distance $\mathcal{W}_{p}^{p}(P,R)$ from the input measure
$P$

$\inf_{R\in\mathcal{P}_{p}(\mathbb{R}^{d})}\left\\{\mathcal{F}(R)+\theta\cdot\mathcal{W}_{p}^{p}(P,R)\right\\}$
(1)

where $\theta>0$ is a weighting parameter. We define the Wasserstein proximal
operator
$\textbf{prox}_{\theta,\mathcal{F}}^{W_{p}}:\mathcal{P}_{p}(\mathbb{R}^{d})\to\mathcal{P}_{p}(\mathbb{R}^{d})$
to be the mapping from input measure $P$ to the optimizer of the (1), $R^{*}$.
That is,

$R^{*}=\textbf{prox}_{\theta,\mathcal{F}}^{\mathcal{W}_{p}}(P):=\arg\inf_{R\in\mathcal{P}_{p}(\mathbb{R}^{d})}\left\\{\mathcal{F}(R)+\theta\cdot\mathcal{W}_{p}^{p}(P,R)\right\\}.$
(2)

The Wasserstein proximal operators for $\mathcal{P}_{p}(\mathbb{R}^{d})$, are
the analogues of proximal operators that arise in convex optimization on
$\mathbb{R}^{d}$ [15] 111Let $f:\mathbb{R}^{d}\to\mathbb{R}$ be a
semicontinuous convex function. The proximal operator of $f$ is a
generalization of a single gradient flow step for $f$ using a backward Euler
scheme with stepsize $1/\lambda$. For $\mathbb{R}^{d}$,
$\textbf{prox}_{\lambda,f}(v)=\arg\min_{x\in\mathbb{R}^{d}}\left[f(x)+\frac{\lambda}{2}\|x-v\|_{2}^{2}\right]$.
Proximal operators are useful for optimizing discontinuous or non-
differentiable functions $f$. . In finite dimensional optimization, proximal
operators are useful for developing optimization methods for nonsmooth or
discontinuous functions. Analogously, the _Wasserstein_ proximal operators are
used to optimize nonsmooth functionals over Wasserstein-$p$ space.

We first review how the Wasserstein-2 and Wasserstein-1 proximals are defined
and implemented in existing generative models, and how the composition of the
proximals is defined.

###  2.1 Wasserstein-2 proximal stabilizes training of generative flows

In canonical continuous normalizing flows, we find a velocity field $v$ such
that the differential equation $\frac{dx}{dt}=v(x(t),t)$ flows a reference
measure $\rho(\cdot,0)=\rho_{0}$ to a target measure $\pi(\cdot)$ given only
samples of the latter. The velocity field is learned by optimizing a
divergence, such as the KL divergence so that the terminal velocity field
$\rho_{T}=\rho(\cdot,T)$ well-approximates $\pi$. In this case, the cost
function is defined
$\mathcal{F}(\rho_{T})=\mathcal{D}_{KL}(\rho_{T}\|\pi)=\mathbb{E}_{\rho_{T}}\left[\log\frac{d\rho_{T}}{d\pi}\right]$.
The Wasserstein-2 proximal of $\mathcal{F}$ is then

$\displaystyle\inf_{\rho_{T}}\left\\{\mathcal{F}(\rho_{T})+\frac{\lambda}{2T}\mathcal{W}_{2}^{2}(\rho_{0},\rho_{T})\right\\}.$
(3)

In relation to (2), here $\theta=\lambda/(2T)$. By the Benamou-Brenier dynamic
formulation of optimal transport, the continuous time optimization problem of
(3) is

$\displaystyle\inf_{\rho_{T},v}\left\\{\mathcal{F}(\rho_{T})+\lambda\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)dxdt:\partial_{t}\rho+\nabla\cdot(v\rho)=0,\,\rho(x,0)=\rho_{0}(x)\right\\}.$
(4)

This is precisely the optimization problem considered in [4] and [3] for
regularizing the training of continuous normalizing flows. The optimal
transport cost can be interpreted as a kinetic energy which favors straighter
paths and makes the optimization more well-posed as it eliminates many valid
flows that may be difficult to integrate. In fact, in Theorem 4.1 we show that
the optimal trajectories are _exactly linear_.

Previous computational studies [3, 4] showed that optimal transport
regularizations stabilized training of CNFs and significantly reduced the
training time. Learning distributions supported on low-dimensional manifolds,
however, is not straightforward, and previous work implemented specialized
architectures or preprocessed the data with autoencoders.

###  2.2 Wasserstein-1 proximal enables manifold learning

The $f$-divergence between two distributions $\rho$ and $\pi$ is defined
$D_{f}(\rho\|\pi)=\mathbb{E}_{\pi}\left[f\left(\frac{d\rho}{d\pi}\right)\right]$,
where $\frac{d\rho}{d\pi}$ is the Radon-Nikodym derivative. A key weakness of
$f$-divergences is that they are only finite when the likelihood ratio
$d\rho/d\pi$ is finite, meaning that they are unable to effectively learn
discrepancies when $\rho$ and $\pi$ are not supported on the same manifold. To
address this deficiency, [7] studied a broad class of regularized
$f$-divergences called the $(f,\Gamma)$-divergence, denoted $D_{f}^{\Gamma}$,
defined via the infimal convolution of the $f$-divergence with an integral
probability metric (IPM) over function space $\Gamma$. Define
$\mathcal{W}_{\Gamma}$ as a general integral probability metric (IPM) over
function space $\Gamma$,
$\mathcal{W}_{\Gamma}(\rho,\nu)=\sup_{\phi\in\Gamma}\left\\{\mathbb{E}_{\rho}[\phi]-\mathbb{E}_{\nu}[\phi]\right\\}.$
The main benefit of IPMs is that they are able to compare distributions with
different supports. The $D_{f}^{\Gamma}$ divergence with weighting parameter
$L>0$ is

$\displaystyle
D_{f}^{\Gamma}(\rho\|\pi)=\inf_{\nu\in\mathcal{P}_{1}(\mathbb{R}^{d})}\left\\{D_{f}(\nu\|\pi)+L\cdot\mathcal{W}_{\Gamma}(\rho,\nu)\right\\}.$
(5)

Via the dual formulation, the Wasserstein-1 distance is an IPM with $\Gamma$
being the space of $1$-Lipschitz functions. Wasserstein-1 proximals of $D_{f}$
are a special $(f,\Gamma)$-divergence called _Lipschitz-regularized_
$f$-divergences, whose primal and dual formulations are

$\displaystyle D_{f}^{\Gamma_{L}}(\rho\|\pi)$
$\displaystyle=\inf_{\nu\in\mathcal{P}_{1}(\mathbb{R}^{d})}\left\\{D_{f}(\nu\|\pi)+L\cdot\mathcal{W}_{1}(\rho,\nu)\right\\}$
(6)
$\displaystyle=\sup_{\phi\in\Gamma_{L}}\left\\{\mathbb{E}_{\rho}[\phi]-\mathbb{E}_{\pi}[f^{\star}(\phi)]\right\\},$
(7)

where $\Gamma_{L}$ is the space of $L$-Lipschitz functions and $f^{\star}$ is
the convex conjugate of $f$. Lipschitz-regularized $f$-divergences are also
related to the so-called Moreau-Yosida $f$-divergences [16].

#### Manifold detecting properties.

By definition of its primal formulation (6), the new divergence satisfies the
following fundamental inequality

$0\leq D_{f}^{\Gamma_{L}}(\rho\|\pi)\leq\min\left\\{D_{f}(\rho\|\pi),L\cdot
W_{1}(\rho,\pi)\right\\}\,.$ (8)

The inequality implies that, unlike $f$-divergences, the new divergence
$D_{f}^{\Gamma_{L}}(\rho\|\pi)$ is able to compare measures $\rho,\pi$ that
are mutually singular. This is especially important when we approximate a
target distribution $\pi$ by the generative distribution $\rho$. When $\pi$ is
supported on an unknown manifold, which is typical in high-dimensional
datasets such as images, the generative distribution $\rho$ can easily be
singular with respect to $\pi$.

#### Lipschitz-regularized $f$-divergences are smooth.

Another key property of $D_{f}^{\Gamma_{L}}(\rho\|\pi)$ is _the existence of a
variational derivative_ with respect to $\rho$, for any
$\rho,\pi\in\mathcal{P}_{1}$, without any absolute continuity or density
assumptions on either $\rho$ or $\pi$ [8]. In addition, the optimizer
$\phi^{*}$ of (7) is the variational derivative of the Lipschitz regularized
$f$-divergence

$\frac{\delta
D_{f}^{\Gamma_{L}}(\rho\|\pi)}{\delta\rho}=\phi^{*}=\arg\max_{\phi\in\Gamma_{L}}\left\\{\mathbb{E}_{\rho}[\phi]-\mathbb{E}_{\pi}[f^{\star}(\phi)]\right\\}\,,$
(9)

A priori the maximizer on the r.h.s of (9) is defined uniquely, up to an
additive constant, in the support of $\rho$ and $\pi$ but the varational
derivative should be understood as the $L$-Lipschitz regularization of the
optimizer which is then defined and $L$-Lipschitz continuous everywhere in
$\mathbb{R}^{d}$, see Lemma 2.3 and Remark 2.4 in [8]. In the context of
adversarial training, $\phi^{*}$ can be interpreted as the _discriminator_ [8,
14]. This fact was first proved in [8] (Theorem 2.1), and was used to
construct data-driven Wasserstein gradient flows of Lipschitz-regularized
$f$-divergences.

The existence of a unique well-defined variational derivative for any
perturbations of all measures in $\mathcal{P}_{1}$ implies _robust_
distribution learning even for distributions supported on a low-dimensional
manifold. Gradient-based optimization methods implicitly approximate the
variational derivative of the cost functional. IPMs, such as
$\mathcal{W}_{1}$, do not have variational derivatives, while $f$\-
divergences cannot handle distributions with differing supports. _The
Wasserstein-1 proximal regularization of $f$-divergences inherits the
beneficial qualities of the $\mathcal{W}_{1}$ distance and $f$-divergences._

##  3 Combining Wasserstein-1 and Wasserstein-2 proximals

The main contribution of this paper is the use of both $\mathcal{W}_{1}$ and
$\mathcal{W}_{2}$ proximals of $f$-divergences to train normalizing flows that
can learn distributions supported on manifolds. The $\mathcal{W}_{1}$ proximal
of $f$-divergences produces Lipschitz-regularized $f$-divergences
$D_{f}^{\Gamma_{L}}$ that enables comparing distributions with differing
supports. The $\mathcal{W}_{2}$ proximal of $D_{f}^{\Gamma_{L}}$ stabilizes
the training of the generative flow by making the optimization problem well-
posed. Let $\rho_{0}$ be the reference distribution, $\pi$ be the target
distribution and $\rho_{T}$ be the approximating distribution to $\pi$. The
combined Wasserstein-1 and Wasserstein-2 proximal of $f$-divergences is
defined as

$\displaystyle\inf_{\rho_{T}}\left\\{D_{f}^{\Gamma_{L}}(\rho_{T}\|\pi)+\frac{\lambda}{2T}\cdot\mathcal{W}_{2}^{2}(\rho_{0},\rho_{T})\right\\}$
(10) $\displaystyle=$
$\displaystyle\underbrace{\inf_{\rho_{T}}\Bigg{\\{}\underbrace{\inf_{\sigma}\left\\{D_{f}(\sigma\|\pi)+L\cdot\mathcal{W}_{1}(\rho_{T},\sigma)\right\\}}_{D_{f}^{\Gamma_{L}}=\mathcal{W}_{1}\text{
proximal of
}D_{f}}+\frac{\lambda}{2T}\cdot\mathcal{W}_{2}^{2}(\rho_{0},\rho_{T})\Bigg{\\}}}_{\mathcal{W}_{2}\text{
proximal of }D_{f}^{\Gamma_{L}}}\,.$ (11)

As is, it is not clear how one may practically compute the objective
functional (11). We describe a computable form of this optimization problem if
we consider optimizing for a continuous time flow. The $\mathcal{W}_{2}$
distance is replaced with the Benamou-Brenier formulation (4). Furthermore, to
compute $D_{f}^{\Gamma}$, we use its dual formulation (7) as it is implemented
in [7, 8]. The resulting optimization problem is

$\displaystyle\inf_{v,\rho}\Bigg{\\{}\overbrace{\sup_{\phi\in\Gamma_{L}}\left\\{\mathbb{E}_{\rho(\cdot,T)}[\phi]-\mathbb{E}_{\pi}[f^{\star}(\phi)]\right\\}}^{\text{Dual
formulation of $D_{f}^{\Gamma_{L}}$
}}+\overbrace{\lambda\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)dxdt}^{\text{Benamou-
Brenier dynamic OT formulation}}\Bigg{\\}}$ (12)
$\displaystyle\frac{dx}{dt}=v(x(t),t),\,x(0)\sim\rho_{0},\,t\in[0,T].$

Again, note that $\Gamma_{L}$ is the space of $L$-Lipschitz continuous
functions and $f^{\star}$ is the convex conjugate of $f$. Roughly speaking,
given a proposed velocity field and samples from $\pi$, the objective
functional (19) is computed by running an ensemble of trajectories with
initial conditions starting from reference distribution $\rho_{0}$. The
optimal transport cost is approximated using trajectories. The final location
of the ensemble determines $\rho(\cdot,T)$, which is compared with samples
from $\pi$ by optimization for a $\phi$ given by the dual formulation of
$D_{f}^{\Gamma}$. In Section 5, we will outline the algorithm for computing
the objective function (19) in detail. We call the resulting flows trained on
this optimization problem $\boldsymbol{\mathcal{W}_{1}\oplus\mathcal{W}_{2}}$
generative flows.

Through a generative adversarial network perspective, we call the optimal
$\phi^{*}$, defined in (9), the _discriminator_ that evaluates the similarity
between the flow-_generated_ distribution $\rho(\cdot,T)$ and the target
distribution $\pi$. The discriminator exists uniquely up to an additive
constant, is well-defined and $L$-Lipschitz on all of $\mathbb{R}^{d}$.

We describe further properties of the discriminator $\phi^{*}$ in Remark 3.1,
and we refer the reader to [7, 8] for further technical details. To connect
with the notion of proximal operators in Wasserstein space (2) [15, 10], we
note that the reference distribution $\rho_{0}$ is related to the generated
distribution $\rho_{T}$ by

$\displaystyle\rho_{T}$
$\displaystyle=\textbf{prox}_{\frac{\lambda}{2T},D_{f}^{\Gamma_{L}}(\cdot\|\pi)}^{\mathcal{W}_{2}}(\rho_{0})=\arg\min_{\nu\in\mathcal{P}_{2}(\mathbb{R}^{d})}\left\\{D_{f}^{\Gamma_{L}}(\nu\|\pi)+\frac{\lambda}{2T}\cdot\mathcal{W}_{2}^{2}(\rho_{0},\nu)\right\\}.$
(13)

######  Remark 3.1 (Composition of Wasserstein proximals and its
interpretation).

One may notice that while we use two Wasserstein proximals to construct the
$\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative flow, there is only a single
($\mathcal{W}_{2}$) proximal operator that relates $\rho_{T}$ to $\rho_{0}$ in
(13). There is an implicit second $\mathcal{W}_{1}$ proximal operator being
computed when calculating the $D_{f}^{\Gamma}$ divergence. In addition to
$\rho_{0}$, $\rho_{T}$, and $\pi$, there is a fourth distribution $\sigma^{*}$
that is produced when comparing $\rho_{T}$ with $\pi$. We have

$\displaystyle\sigma^{*}=\textbf{prox}_{L,D_{f}(\cdot\|\pi)}^{\mathcal{W}_{1}}(\rho_{T})=\arg\min_{\sigma\in\mathcal{P}_{1}(\mathbb{R}^{d})}\left\\{D_{f}(\sigma\|\pi)+L\cdot\mathcal{W}_{1}(\rho_{T},\sigma)\right\\},$
(14)

where $\sigma^{*}$ can be interpreted as an intermediate measure that is close
to $\rho_{T}$ in the $\mathcal{W}_{1}$ distance while sharing the same support
as $\pi$ so that they can be compared via the $f$-divergence. Moreover, by
(13), we see that $\sigma^{*}$ is equal to a _composition_ of the
$\mathcal{W}_{1}$ and $\mathcal{W}_{2}$ proximal operators

$\displaystyle\sigma^{*}=\left(\textbf{prox}_{L,D_{f}(\cdot\|\pi)}^{\mathcal{W}_{1}}\circ\textbf{prox}_{\frac{\lambda}{2T},D_{f}^{\Gamma_{L}}(\cdot\|\pi)}^{\mathcal{W}_{2}}\right)(\rho_{0}).$
(15)

The measure $\sigma^{*}$ is not explicitly computed in our algorithm as we opt
to work with the dual formulation of the $\mathcal{W}_{1}$ proximal of
$D_{f}$. However, $\sigma^{*}$ is related to the optimal discriminator
$\phi^{*}$ by the relation

$\displaystyle\frac{d\sigma^{*}}{d\pi}(x)=(f^{\star})^{\prime}(\phi^{*}(x)-c)$
(16)

for some constant $c$. That is, the likelihood ratio of $\sigma^{*}$ with
respect to the target distribution $\pi$ is equal to the derivative of
$f^{\star}$ evaluated at a constant shift of the optimal discriminator
$\phi^{*}$. See [7], Theorem 25 for further details.

##  4 $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative flows as mean-field
games

In this section, we study our new $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$
generative flow as a mean-field game (MFG) and discuss the conditions for the
combined proximal training objective to be _well-posed_. We will see that the
well-posedness of our optimization can be understood through MFG theory. After
reviewing the basics of MFG theory for generative flows, we prove that the
$\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative flows have a well-defined
PDE system even for high-dimensional data supported on manifolds. The PDE
solutions are _unique_ assuming that the solution is smooth, which implies the
optimization problem is well-posed, and yields a _unique optimal generative
flow_. This uniqueness property of the generative flow is expected to render
corresponding algorithms more stable and not subject to implementation-
dependent outcomes. Moreover, we show that the optimal trajectories are linear
through connections to PDE theory and fluid mechanics.

###  4.1 Background on continuous-time generative flows as mean-field games

Generative flows were shown to be trained by solving mean-field game problems
in [9]. In a general sense, the mean-field game for deterministic generative
flows is given by a terminal cost functional
$\mathcal{F}:\mathcal{P}(\mathbb{R}^{d})\to\mathbb{R}$, such as the KL
divergence with respect to target distribution $\pi$, which measures the
discrepancy between $\pi$ and approximating distribution $\rho(\cdot,T)$, and
a convex running cost
$L(x,v):\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$, which measures the
cost particle moving at a particular velocity $v$ at position $x$. The mean-
field game is

$\displaystyle\begin{split}&\min_{v,\rho}\,\mathcal{F}(\rho(\cdot,T))+\int_{0}^{T}\int_{\mathbb{R}^{d}}L(x,v(x,t))\rho(x,t)dxdt\\\
\text{s.t.
}&\partial_{t}\rho+\nabla\cdot(v\rho)=0\,\,\,\,\rho(x,0)=\rho_{0}(x)\end{split}$
(17)

Via a Lagrangian interpretation, the continuity equation can be replaced by
particle trajectories that the differential equation $dx/dt=v(x(t),t)$, where
$x(t)\sim\rho(\cdot,t)$. The advantage of the MFG perspective is that its
associated _optimality conditions_ , which are in the form of a system of
backward-forward PDEs, describe the mathematical structure of the optimal
velocity field and density evolution _a priori_ to computation. Moreover, the
lack of optimality conditions implies _ill-posedness_ of the optimization
problem.

Define the Hamiltonian of the MFG to be the convex conjugate of the running
cost function $L$, $H(x,p)=\sup_{v}\left[-p^{\top}v-L(x,v)\right]$. The
optimizers of (17) are characterized by the optimality conditions

$\displaystyle\begin{dcases}-\frac{\partial U}{\partial t}+H(x,\nabla
U)=0,\,\,\,&U(x,T)=\frac{\delta\mathcal{F}(\rho(\cdot,T))}{\delta\rho(\cdot,T)}(x),\\\
\frac{\partial\rho}{\partial t}-\nabla\cdot(\nabla_{p}H(x,\nabla
U)\rho)=0,\,\,\,\,\,&\rho(x,0)=\rho_{0}(x).\end{dcases}$ (18)

The first equation is a backward Hamilton-Jacobi (HJ) equation, while the
second equation is the forward continuity equation with the optimal velocity
field is determined by the solution to the HJ equation
$v^{*}=-\nabla_{p}H(x,\nabla U)$. The well-posedness of the optimization
problem for training CNFs (17) can then be studied via PDE theory of the
optimality conditions (18).

Empirically, the optimization problem for learning canonical CNFs, i.e., when
$\mathcal{F}(\rho)$ is the KL divergence and $L(x,v)=0$, is known to exhibit
instabilities during training as there are multiple flows that can minimize
the objective functional [5, 3, 4]. This fact can be immediately seen through
MFGs since with $L=0$, there is no well-defined Hamiltonian, meaning there is
no HJ equation, and therefore no optimality conditions. It is also easy to see
that there are infinitely many optimal flows in this setting. See [9] for
further details. On the other hand, with an optimal transport cost
$L(x,v)=\frac{1}{2}|v|^{2}$, [4, 3] found that training stabilized, and [4]
derived the associated Hamilton-Jacobi equation and used it to further
accelerate the training process.

###  4.2 MFG for $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative flows have
well-defined optimality conditions and imply _linear_ optimal trajectories

Wasserstein-2 proximals stabilize the training of generative flows, while
Wasserstein-1 proximals enables manifold learning. In this section, we back up
these claims by studying the PDE system that arises from their optimality
conditions. In particular, we will see the $\mathcal{W}_{2}$ proximal yields a
well-defined Hamilton-Jacobi equation for a potential function $U(x,t)$, while
the $\mathcal{W}_{1}$ proximal produces well-defined terminal condition
$U(x,T)$.

######  Theorem 4.1 (MFG for $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ proximal
generative flow).

Let $\pi$ be an unknown arbitrary target measure, $\rho(x,0)=\rho_{0}$ be a
given reference measure, and
$v:\mathbb{R}^{d}\times\mathbb{R}\rightarrow\mathbb{R}^{d}$ be a vector field.
Fix a terminal time $T>0$ and $\lambda>0$. The optimization problem

$\inf_{v,\rho}\left\\{\sup_{\phi\in\Gamma_{L}}\left\\{\mathbb{E}_{\rho(\cdot,T)}[\phi]-\mathbb{E}_{\pi}[f^{\star}(\phi)]\right\\}+\lambda\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)dxdt\right\\}$
(19)

where $\rho(x,t)$ satisfies the continuity equation
$\partial_{t}\rho+\nabla\cdot(\rho v)=0,$ $\rho(x,0)=\rho_{0}(x)$ has the
following optimality conditions:

  1. 1.

Wasserstein-2 proximal yields a well-posed Hamilton-Jacobi equation.  The
optimal velocity $v^{*}$ is given by $v^{*}=-\frac{1}{\lambda}\nabla U$ where
$U$ and $\rho$ satisfy a system consisting of a backward Hamilton-Jacobi
equation and a forward continuity equation for $t\in[0,T]$

$\displaystyle\begin{dcases}-\partial_{t}U+\frac{1}{2\lambda}|\nabla
U|^{2}=0\,\,\,\,\,&\\\ \partial_{t}\rho-\nabla\cdot\left(\rho\frac{\nabla
U}{\lambda}\right)=0\,\,\,\,\,&\text{where}\,\,\,\,\rho(x,0)=\rho_{0}(x).\end{dcases}$
(20)

  2. 2.

Wasserstein-1 proximal provides a well-defined terminal condition.  The dual
formulation and the variational derivative of the Wasserstein-1 proximal
$D_{f}^{\Gamma_{L}}(\rho(\cdot,T)\|\pi)$, (7) and (9) respectively, provide
the terminal condition of the Hamilton-Jacobi equation given as

$U(x,T)=\frac{\delta
D_{f}^{\Gamma_{L}}(\rho(\cdot,T)\|\pi)}{\delta\rho(\cdot,T)}(x)=\arg\max_{\phi\in\Gamma_{L}}\left\\{\mathbb{E}_{\rho(\cdot,T)}[\phi]-\mathbb{E}_{\pi}[f^{\star}(\phi)]\right\\}=\phi^{*}(x).$
(21)

  3. 3.

Optimal trajectories are linear. The optimal velocity field is
$v(x,t)=-\frac{1}{\lambda}\nabla U(x,t)$. For an initial condition
$x(0)\sim\rho(\cdot,0)$, the optimal trajectories have the following
representations

$\displaystyle x(t)$ $\displaystyle=x(0)-\frac{1}{\lambda}\int_{0}^{t}\nabla
U(x(s),s)ds=x(0)-\frac{t}{\lambda}\nabla U(x(0),0)$
$\displaystyle=x(T)-\frac{1}{\lambda}\int_{T}^{t}\nabla
U(x(s),s)ds=x(T)+\frac{T-t}{\lambda}\nabla U(x(T),T)$ (22)
$\displaystyle=x(T)+\frac{T-t}{\lambda}\nabla\phi^{*}(x(T))$

meaning that they are _exactly_ linear.

###### Proof.

The proof is a direct application of the general MFG optimality conditions in
(18) [9]. Observe that the Hamiltonian is

$\displaystyle
H(x,p)=\sup_{v\in\mathbb{R}^{d}}\left\\{-p^{\top}v-\frac{\lambda}{2}|v|^{2}\right\\}=\frac{1}{2\lambda}|p|^{2}.$
(23)

A key element of this proof is the use the variational derivative (9) of the
Wasserstein-1 proximal $D_{f}^{\Gamma_{L}}(\rho(\cdot,T)\|\pi)$, which exists
for any probability measures in $\mathcal{P}_{1}$ without any absolute
continuity or density assumptions, and is equal to the discriminator
$\phi^{*}$, as discussed in Section 2.2 and [8]. This terminal condition is
$L$-Lipschitz continuous everywhere in $\mathbb{R}^{d}$.

To see that the optimal trajectories are linear, consider the (system of) PDEs
associated with the optimal velocity field $v(x,t)=-\frac{1}{\lambda}\nabla
U(x,t)$, which is derived from the Hamilton-Jacobi equation

$\displaystyle-\frac{1}{\lambda}\partial_{t}\nabla
U+\frac{1}{\lambda^{2}}\nabla U\cdot\nabla(\nabla
U)=0\implies\partial_{t}v+(v\cdot\nabla)v=0.$ (24)

We implicitly assume that $U$ is sufficiently differentiable. See Remark 4.1
for further comments on regularity. Note that together with the continuity
equation in (20), this system of PDEs arises in inviscid compressible fluid
mechanics and is called the Euler equations [17]. Next, consider the
_characteristics_ for the PDEs of $v$, i.e., the level curves on which $v$ is
constant [18]. Let $y(t)$ be the characteristics which satisfy

$\displaystyle\frac{d}{dt}v(y(t),t)=\partial_{t}v(y(t),t)+\partial_{t}y(t)\cdot\nabla
v(y(t),t)=0.$ (25)

However, from (24), we know that $y(t)$ is a characteristic only if

$\displaystyle\partial_{t}y(t)=v(y(t),t)=-\frac{1}{\lambda}\nabla U(y(t),t)$
(26)

which is exactly the evolution of the optimal trajectories $x(t)$ in (3). This
means that the optimal trajectories are precisely the characteristics of the
velocity field PDEs, implying that the _velocity is constant along the optimal
trajectories_. Therefore, $v(x(t),t)=v(x(0),0)=-\frac{1}{\lambda}\nabla
U(x(0),0)$ for all $t\in[0,T]$, and so

$\displaystyle x(t)=x(0)+\int_{0}^{t}v(x(s),s)ds=x(0)-\frac{t}{\lambda}\nabla
U(x(0),0).$ (27)

The optimal trajectories may instead be characterized by with their terminal
endpoints. Along optimal trajectories,
$v(x(t),t)=v(x(T),T)=-\frac{1}{\lambda}\nabla U(x(T),T)$. The terminal
condition is equal to the discriminator, i.e., $U(x,T)=\phi^{*}(x)$.
Therefore,

$\displaystyle x(t)$ $\displaystyle=x(T)-\frac{1}{\lambda}\int_{T}^{t}\nabla
U(x(s),s)ds=x(T)+\frac{T-t}{\lambda}\nabla U(x(T),T)$ (28)
$\displaystyle=x(T)+\frac{T-t}{\lambda}\phi^{*}(x(T)).$ (29)

∎

###  4.3 Wasserstein-1 proximal provides robust learning of distributions on
manifolds

The selection of the Wasserstein-1 proximal divergence (6) and (7) as terminal
condition in (19) relied on two important features of this new divergence
$D_{f}^{\Gamma_{L}}(P\|Q)$: (a) it can compare mutually singular measures—a
property inherited by Wasserstein-1—as can be readily seen in (8); (b) it is
smooth, i.e. has a variational derivative (9), a property inherited by the KL
divergence or other f-divergence. The latter is important in the well-
posedness of the generative flow as captured by the MFG optimality conditions
(20) and (21) of Theorem 4.1.

####  $f$-divergences fail to learn manifolds.

When an $f$-divergence is used as the terminal cost to train generative flows,
even with a $\mathcal{W}_{2}$ proximal, the mean-field game (19) will not be
well-posed. The terminal condition (21) is the variational derivative of
$\mathcal{F}(\rho)=D_{f}(\rho\|\pi)$ for a fixed $\pi$, $\frac{\delta
D_{f}(\rho\|\pi)}{\delta\rho}=(f^{\star})^{\prime}\left(\frac{d\rho}{d\pi}\right)$.
In particular, for the KL divergence

$\displaystyle\frac{\delta
D_{KL}(\rho\|\pi)}{\delta\rho}(x)=1+\log\frac{d\rho}{d\pi}(x).$ (30)

If $\rho$ is defined as the pushforward of a standard normal distribution in
$\mathbb{R}^{d}$ under invertible mapping, and the support of $\pi$ is on a
lower dimensional manifold, then the terminal condition will blow up for
$x\notin\text{supp}(\pi)$. In essence, the KL divergence becomes
_uninformative_ for learning, when the target distribution is supported on a
different manifold than $\rho$.

####  $\mathcal{W}_{1}$ proximal of $D_{f}$ provides robust manifold learning

First, (8) implies that the Wasserstein-1 proximal regularized $f$-divergence
ball, $\\{\rho:D_{f}^{\Gamma_{L}}(\rho\|\pi)\leq 1\\}$ contains both the
$f$-divergence ball $\\{\rho:D_{f}^{\Gamma_{L}}(\rho\|\pi)\leq 1\\}$ and the
Wasserstein-1 ball $\\{\rho:L\cdot W_{1}(\rho,\pi)\leq 1\\}$. This is an
important point in our context because algorithmically a KL or an
$f$-divergence ball around a target distribution $\pi$ that is supported on a
low-dimensional manifold, as is typical in many high-dimensional data sets
such as images, will only include other distributions with supports on the
_same_ (unknown) manifold. On the other hand a ball in
$D_{f}^{\Gamma_{L}}(\cdot\|\pi)$ does not have this constraint and hence it is
expected to be mathematically and computationally more robust for generative
tasks. In particular, when learning a target distribution $\pi$ supported on
an unknown manifold, generated models $\rho_{T}$ can easily be singular with
respect to $\pi$ unless great care is employed using carefully constructed
autoencoders and/or neural architectures. Our use of the Wasserstein-1
proximal $D_{f}^{\Gamma_{L}}(\cdot\|\pi)$ does away with such needs, as
Theorem 4.1 suggests and as we demonstrate computationally in Section 6. This
discussion also has connection to distributionally robust optimization (DRO)
[19].

#### Well-defined variational derivative implies stable training

The variational derivative $\phi^{*}$ in (9) for the terminal condition exists
and is always well-defined due to our choice of the terminal cost
$\mathcal{F}$ as the Wasserstein-1 proximal of an $f$-divergence
$D_{f}^{\Gamma_{L}}(\cdot\|\pi)$. Here the use of a proximal is necessary: (i)
in the case of a pure $f$-divergence the variational derivative will blow up
for mutually singular measures, while (ii) in the case of a pure Wasserstein-1
distance, there is no uniquely defined variational derivative. The MFG well-
posedness analysis of Theorem 4.1 allows us to make rigorous the following
heuristics. In the optimization process over possible velocities $v=v(x,t)$ in
our learning objective in Theorem 4.1, namely

$\displaystyle\inf_{v,\rho}$
$\displaystyle\left\\{D_{f}^{\Gamma_{L}}(\rho_{T}\|\pi)+\lambda\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)dxdt\right\\}\,,$
(31)

the existence of the variational derivative (9) of the terminal cost
$D_{f}^{\Gamma_{L}}(\rho_{T}\|\pi)$ for any perturbations of measures
including singular ones, suggests that this minimization problem is smooth in
$v$ and the optimal generative flow can be theoretically discovered by
gradient optimization. We discuss and implement such algorithms inspired by
Theorem 4.1 in Section 5.

###  4.4 Uniqueness of $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative flows
implies well-posedness of optimization problem

In Theorem 4.1 we have shown that the optimization problem (19) and the
corresponding backward-forward MFG system (20) are well-defined in the sense
that the Hamiltonian $H$ exists and the terminal condition (21) is always
well-defined. Next we show that (20) and (21) have a unique smooth solution
for a bounded domain with periodic boundary conditions, i.e., a torus. This is
a key result of this paper because it implies that the corresponding
generative flow with velocity field $v^{*}(x,t)$ is also unique and therefore
can be learned in a stable manner. We will also demonstrate that latter point
in our experiments in Section 6.

######  Theorem 4.2 (Uniqueness of Wasserstein-1/Wasserstein-2 proximal
generative flows).

If the backward-forward PDE system (20) with terminal condition (21) has
smooth solutions $(U,\rho^{*})$ on the torus $\Omega$, then they are unique
and the solution to the optimization problem (19) is also unique.

We follow the strategy outlined in Theorem 2.5 of [20]. The main idea of the
proof is to consider the weak form of the continuity equation where we use the
solutions of the HJ equation as test functions. We focus on the periodic
domain (torus) case for the sake of simplicity, although other boundary
conditions can be considered, see [20]. The proof requires two properties: (a)
the convexity of the Hamiltonian $H(x,p)=\frac{1}{2\lambda}|p|^{2}$; and (b)
the terminal conditions (21) is necessarily increasing in some suitable sense
in $\rho$. Convexity of the Hamiltonian is clear. As for the mononicity
requirement, we first prove the following Lemma 4.1. The result is based on
the properties of the $\mathcal{W}_{1}$-proximal.

######  Lemma 4.1 (Monotonicity of (21)).

If $\phi_{i}^{*}$ is the optimizer (21) of the dual problem (7) for
$D_{f}(\rho_{i}(\cdot,T)\|\pi)$ for $i=1,2$, then

$\int_{\Omega}(\phi_{1}^{*}(x)-\phi_{2}^{*}(x))d(\rho_{1}(\cdot,T)-\rho_{2}(\cdot,T))(x)\geq
0$ (32)

for all $\rho_{1}(\cdot,T),\rho_{2}(\cdot,T)\in\mathcal{P}(\Omega)$.

###### Proof.

For the simplicity we write $\mu_{1}=\rho_{1}(\cdot,T)$ and
$\mu_{2}=\rho_{2}(\cdot,T)$. We rewrite (32) and use the optimality of
$\phi_{i}^{*}$ for the dual formulation (7) of
$D_{f}^{\Gamma_{L}}(\mu_{i}\|\pi)$ with $i=1,2$ to get

$\displaystyle\begin{split}&\mathbb{E}_{\mu_{1}}[\phi_{1}^{*}]-\mathbb{E}_{\mu_{2}}[\phi_{1}^{*}]-\mathbb{E}_{\mu_{1}}[\phi_{2}^{*}]+\mathbb{E}_{\mu_{2}}[\phi_{2}^{*}]\\\
=&\\{D_{f}^{\Gamma_{L}}(\mu_{1}\|\pi)+\mathbb{E}_{\pi}[f^{\star}(\phi_{1}^{*})]\\}-\mathbb{E}_{\mu_{2}}[\phi_{1}^{*}]-\mathbb{E}_{\mu_{1}}[\phi_{2}^{*}]+\\{\mathbb{E}_{\pi}[f^{\star}(\phi_{2}^{*})]+D_{f}^{\Gamma_{L}}(\mu_{2}\|\pi)\\}\\\
\geq&D_{f}^{\Gamma_{L}}(\mu_{1}\|\pi)-D_{f}^{\Gamma_{L}}(\mu_{2}\|\pi)-D_{f}^{\Gamma_{L}}(\mu_{1}\|\pi)+D_{f}^{\Gamma_{L}}(\mu_{2}\|\pi)\geq
0.\end{split}$ (33)

∎

We now prove Theorem 4.2.

###### Proof.

Let $(U_{1},\rho_{1})$ and $(U_{2},\rho_{2})$ be two smooth solutions of (20)
and (21). Define $U=U_{1}-U_{2}$ and $\rho=\rho_{1}-\rho_{2}$. We prove $U=0$
and $\rho=0$. Based on Lemma 4.1, we have the monotonicity of the terminal
condition $U(\cdot,T)$. Furthermore, $H(x,p)=\frac{1}{2\lambda}|p|^{2}$, is
convex with respect to the last variable. Without loss of generality, let
$\lambda=1$. Observe that $\nabla_{p}^{2}H(x,p)=I_{d}$. Now observe that

$\displaystyle\begin{split}&\frac{d}{dt}\int_{\Omega}U\rho
dx=\int_{\Omega}[(\partial_{t}U)\rho+U(\partial_{t}\rho)]dx\\\
=&\int_{\Omega}\left(H(x,\nabla U_{1})-H(x,\nabla U_{2})\right)\rho dx\\\
+&\int_{\Omega}U\left(\nabla\cdot(\rho_{1}\nabla_{p}H(x,\nabla
U_{1}))-\nabla\cdot(\rho_{2}\nabla_{p}H(x,\nabla U_{2}))\right)dx.\end{split}$
(34)

Via integration by parts, with $\langle\cdot,\cdot\rangle$ denoting the
standard inner product between vectors in $\mathbb{R}^{d}$, we get

$\displaystyle\begin{split}\eqref{eq:proof:uniqueness:eq1}=&\int_{\Omega}\left[(H(x,\nabla
U_{1})-H(x,\nabla U_{2}))\rho-\langle\nabla U,\rho_{1}\nabla_{p}H(x,\nabla
U_{1})-\rho_{2}\nabla_{p}H(x,\nabla U_{2})\rangle\right]dx\\\
=&-\int_{\Omega}\rho_{1}\left[H(x,\nabla U_{2})-H(x,\nabla
U_{1}))-\langle\nabla U_{2}-\nabla U_{1},\nabla_{p}H(x,\nabla
U_{1})\rangle\right]dx\\\ -&\int_{\Omega}\rho_{2}\left[H(x,\nabla
U_{1})-H(x,\nabla U_{2}))-\langle\nabla U_{1}-\nabla
U_{2},\nabla_{p}H(x,\nabla U_{2})\rangle\right]dx.\end{split}$ (35)

Due to the convexity of $H$, a Taylor series expansion yields

$\eqref{eq:proof:uniqueness:eq2}=\frac{d}{dt}\int_{\Omega}U\rho
dx\leq-\int_{\Omega}\frac{\rho_{1}+\rho_{2}}{2}|\nabla U_{1}-\nabla
U_{2}|^{2}dx\leq 0.$ (36)

By integrating this inequality on the time interval $[0,T]$ we obtain

$\int_{\Omega}U(x,T)\rho(x,T)dx\leq\int_{\Omega}U(x,0)\rho(x,0)dx-\int_{0}^{T}\int_{\Omega}\frac{\rho_{1}+\rho_{2}}{2}|\nabla
U_{1}-\nabla U_{2}|^{2}dxdt$ (37)

Using that $\rho(x,0)=0$ and the monotonicity (32), we have the LHS of (37) is
nonnegative. It implies that $\int_{\Omega}U(x,T)\rho(x,T)dx=0$ resulting in

$\nabla U_{1}=\nabla U_{2}\text{ in }\\{\rho_{1}>0\\}\cup\\{\rho_{2}>0\\}.$
(38)

This proves that $\rho_{1}$ solves the same equation as $\rho_{2}$, with the
same drift $\nabla_{p}H(x,\nabla U_{1})=\nabla_{p}H(x,\nabla U_{2})$ and
therefore $\rho_{1}=\rho_{2}$. Correspondingly $U_{1}$ and $U_{2}$ solve the
same HJ equation and therefore $U_{1}=U_{2}$, as they share the same terminal
condition (21). ∎

From a computational perspective, this theorem provides confidence that any
numerical implementation to the optimization problem (12) is approximating a
single unique limiting solution. Algorithmically, this means that the training
will not oscillate between multiple modes representing different possible
limiting flows.

######  Remark 4.1.

In Theorem 4.2 we assumed the existence of classical solutions to prove
uniqueness. However first order HJ do not always have classical solutions,
hence further investigation is needed in future work. For instance, we intend
to consider _stochastic_ normalizing flows for which the corresponding MFG has
classical solutions [20], and a uniqueness result in the spirit of Theorem 4.2
can hold without additional regularity assumptions.

##  5 Proximal regularization produces adversarial training algorithm

In this section, we provide more details as to how the objective function for
training $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ proximal generative flows are
computed. The mean-field game formulation in (19) enables implementation of
the new flow through _adversarial_ training algorithms. This is due to the
dual formulation for computing the $D_{f}^{\Gamma}$ divergence, and has been
successfully applied in generative modeling tasks through
$D_{f}^{\Gamma}$-based generative adversarial networks [7, 14, 13], and
Wassserstein gradient flows of $D_{f}^{\Gamma}$ [8]. A generative adversarial
flow has previously been proposed in [12]. Our formulation, however, resolves
the persistent ill-posedness issues.

###  5.1 An adversarial training algorithm for
$\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative flows

We compute the objective function in (19), see also (12), and we will see that
in contrast to normalizing flows and its $\mathcal{W}_{2}$ proximal
regularizations [2, 3, 4], we only need to simulate forward trajectories of
the flow during training. Suppose we are given a dataset
$\\{X^{(n)}\\}_{n=1}^{N}\sim\pi$. Let the discriminator
$\phi(x;\theta_{\phi})$ and potential function $U(x,t;\theta_{U})$ be
parametrized via neural networks, with parameters $\theta_{\phi}$ and
$\theta_{U}$, respectively. Recall that $\phi$ needs to be $L$-Lipschitz
continuous. Fix an initial reference distribution $\rho_{0}$ (e.g., a normal
distribution), and simulate $M$ independent trajectories with initial
condition $\\{Y_{0}^{(m)}\\}_{m=1}^{M}\sim\rho_{0}$ and velocity field
$v(y,t)=-\frac{1}{\lambda}\nabla_{x}U(y(t),t)$. The number of simulated
trajectories $M$ do not need to match the number of samples $N$ from $\pi$. We
follow a discretize-then-optimize approach for training the model parameters,
as described in [4]. Each trajectory is simulated using a forward Euler
scheme. Discretize the time interval $[0,T]$ into $K$ equal segments of length
$h$ so that $t_{K}=T=Kh$. Then for the $m$th trajectory at time $t_{k}=kh$, we
have the recurrence relation

$\displaystyle Y_{k+1}^{(m)}=Y_{k}^{(m)}-\frac{h}{\lambda}\nabla
U(Y_{k}^{(m)},kh)\sim\rho(\cdot,(k+1)h),\,\,\,\,\text{ for }k=0,\ldots,K-1.$
(39)

We first estimate the divergence $D_{f}^{\Gamma}(\rho(\cdot,T)\|\pi)$ by
finding the optimal $\phi\in\Gamma_{L}$

$\displaystyle\max_{\theta_{\phi}}\mathcal{J}_{\phi}(\theta_{\phi})=\max_{\theta_{\phi}}\left\\{\frac{1}{M}\sum_{m=1}^{M}\phi\left(Y^{(m)}_{K};\theta_{\phi}\right)-\frac{1}{N}\sum_{n=1}^{N}f^{\star}\left(\phi(X^{(n)};\theta_{\phi})\right)+\mathcal{L}_{\phi}(\theta_{\phi})\right\\}$
(40) $\displaystyle\text{where
}~{}~{}\mathcal{L}_{\phi}(\theta_{\phi})=-\sum_{n=1}^{\min(M,N)}\max\left(\left|\nabla\phi(c_{n}Y_{K}^{(n)}+(1-c_{n})X^{(n)};\theta_{\phi})\right|^{2}-L^{2},0\right).$
(41)

Here, $\mathcal{L}_{\phi}(\theta_{\phi})$ penalizes discriminators that have
Lipschitz constant greater than $L$ for random convex combinations of
$Y_{K}^{(n)}$ and $X^{(n)}$ [7, 21]. Here, $c_{n}$ is randomly sampled from a
uniform distribution on $[0,1]$. Discrminator parameters $\theta_{\phi}$ are
updated according to a gradient ascent method
$\theta_{\phi}^{(l+1)}=\theta_{\phi}^{(l)}+\eta\nabla_{\theta_{\phi}}\mathcal{J}_{\phi}(\theta_{\phi}^{(l)})$
for some learning rate $\eta$. Let $\phi^{*}(x)=\phi(x;\theta^{*}_{\phi})$ the
resulting optimal discriminator. Thereafter, the objective function
$\mathcal{J}(\theta_{U})$ for learning the parameters of $U$ is estimated by

$\displaystyle\mathcal{J}_{U}(\theta_{U})=\underbrace{\frac{1}{M}\sum_{m=1}^{M}\phi^{*}\left(Y^{(m)}_{K}\right)-\frac{1}{N}\sum_{n=1}^{N}f^{\star}\left(\phi^{*}(X^{(n)})\right)}_{\text{Estimates
}D_{f}^{\Gamma_{L}}(\rho(\cdot,T)\|\pi)}+\underbrace{\frac{\lambda
h}{2M}\sum_{m=1}^{M}\sum_{k=0}^{K-1}\left|\frac{\nabla
U(Y_{k}^{(m)},kh;\theta_{U})}{\lambda}\right|^{2}}_{\text{Estimates
}\lambda\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)dxdt}.$
(42)

Model parameters $\theta_{U}$ are similarly updated via gradient descent:
$\theta_{U}^{(l+1)}=\theta_{U}^{(l)}-\eta\nabla_{\theta_{U}}\mathcal{J}_{U}(\theta_{U}^{(l)})$.
From here, new sample trajectories are simulated using the potential at step
$l+1$, $U^{(l+1)}(x,t)=U(x,t;\theta_{U}^{(l+1)})$ and the parameters are
updated again. Like GANs, there is a minimax structure in the optimization
problem, where the inner optimization problem is solved for each update of the
outer problem. Our proposed generative model can be viewed as an _adversarial
generative flow_ given in terms of a minimax problem for a discriminator
$\phi^{*}$ and a potential function $U^{*}$. The pseudocode for learning
$U^{*}$ and $\phi^{*}$ are described in Algorithms 1 and 2, respectively.

Algorithm 1 Training $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ proximal
generative flows

1:Input: Dataset $\\{X^{(n)}\\}_{n=1}^{N}\sim\pi$

2:Initialize: Discriminator $\phi(x;\theta_{\phi})$ and potential function
$U(x,t;\theta_{U})$ with parameters $\theta_{\phi}$ and $\theta_{U}$,
Lipschitz constant $L$, initial reference distribution $\rho_{0}$, number of
trajectories $M$, time interval $[0,T]$, with step size $h$, $T=Kh$, learning
rate $\eta$, $N_{iter}^{\phi}$, $N_{iter}^{U}$

3:for $l\leftarrow 1$ to $N^{U}_{iter}$ do

4: Sample $M$ initial conditions $\\{Y_{0}^{(m)}\\}_{m=1}^{M}\sim\rho_{0}$

5: for $m=1$ to $M$ do $\triangleright$ Simulate trajectories

6: for $k=0$ to $K-1$ do

7: $Y_{k+1}^{(m)}=Y_{k}^{(m)}-\frac{h}{\lambda}\nabla
U(Y_{k}^{(m)},kh;\theta_{U})$

8: end for

9: end for

10: $\phi^{*}(x)\leftarrow$
LearnDiscriminator$(\\{X^{(n)}\\},\\{Y_{K}^{(m)}\\},L,\eta,N_{iter}^{\phi})$

11:
$\mathcal{D}_{U}(\theta_{U})\leftarrow\frac{1}{M}\sum_{m=1}^{M}\phi^{*}(Y^{(m)}_{K})-\frac{1}{N}\sum_{n=1}^{N}f^{\star}(\phi^{*}(X^{(n)}))$
$\triangleright$ Estimate $D_{f}^{\Gamma_{L}}(\rho_{T}\|\pi)$

12:
$\mathcal{J}_{U}(\theta_{U})\leftarrow\mathcal{D}_{U}(\theta_{U})+\frac{\lambda
h}{2M}\sum_{m=1}^{M}\sum_{k=0}^{K-1}\left|\frac{\nabla
U(Y_{k}^{(m)},kh;\theta_{U})}{\lambda}\right|^{2}$ $\triangleright$ Estimate
objective

13:
$\theta_{U}\leftarrow\theta_{U}-\eta\nabla_{\theta_{U}}\mathcal{J}_{U}(\theta_{U})$
$\triangleright$ Update $\theta_{U}$ via gradient descent

14:
$\mathcal{R}^{(l)}_{HJ}\leftarrow\frac{h}{M}\sum_{m=1}^{M}\sum_{k=0}^{K-1}\left|\partial_{t}U(Y_{k}^{(m)},hk;\theta_{U})+\frac{1}{2\lambda}\left|\nabla
U(Y_{k}^{(m)},hk;\theta_{U})\right|^{2}\right|$

15: $\mathcal{R}_{T}^{(l)}\leftarrow\frac{1}{M}\sum_{m=1}^{M}\left|\nabla
U(Y_{K}^{(m)},T;\theta_{U})-\nabla\phi^{*}(Y_{K}^{(m)})\right|$
$\triangleright$ Optimality indicators (43), (44)

16:end for

Algorithm 2 Learning the discriminator $\phi^{*}(x)$

1:function
LearnDiscriminator($\\{X^{(n)}\\}_{n=1}^{N},\\{Y_{K}^{(m)}\\}_{m=1}^{M},L,\eta,N_{iter}^{\phi}$)

2: Initialize: Discriminator $\phi(x;\theta_{\phi})$ with parameters
$\theta_{\phi}$

3: for $l\leftarrow 1$ to $N_{iter}^{\phi}$ do

4: $c_{n}\sim\mathcal{U}[0,1]$ for $n=1,\ldots,\min(M,N)$

5: $Z^{(n)}\leftarrow c_{n}Y_{K}^{(n)}+(1-c_{n})X^{(n)}$ for
$n=1,\ldots,\min(M,N)$

6:
$\mathcal{L}_{\phi}(\theta_{\phi})\leftarrow-\sum_{n=1}^{\min(M,N)}\max\left(\left|\nabla\phi(Z^{(n)};\theta_{\phi})\right|^{2}-L^{2},0\right)$
$\triangleright$ Lipschitz regularization

7:
$\mathcal{D}_{\phi}(\theta_{\phi})\leftarrow\frac{1}{M}\sum_{m=1}^{M}\phi(Y^{(m)}_{K};\theta_{\phi})-\frac{1}{N}\sum_{n=1}^{N}f^{\star}(\phi(X^{(n)};\theta_{\phi}))$

8:
$\mathcal{J}_{\phi}(\theta_{\phi})\leftarrow\mathcal{D}_{\phi}(\theta_{\phi})+\mathcal{L}_{\phi}(\theta_{\phi})$

9:
$\theta_{\phi}\leftarrow\theta_{\phi}+\eta\nabla_{\theta_{\phi}}\mathcal{J}_{\phi}(\theta_{\phi})$
$\triangleright$ Update $\theta_{\phi}$ via gradient ascent

10: end for

11: return $\phi^{*}(x)=\phi(x;\theta_{\phi})$

12:end function

###  5.2 Hamilton-Jacobi optimality indicators

The optimality conditions (20) and (21) of Theorem 4.1 provide _a posteriori_
error estimators that may provide insight into the quality of the learned
generative flow. Recall that the Hamilton-Jacobi equation (43) and its
terminal condition (44) characterize the optimal velocity field of the
generative flow, which means at the end of training, the learned flow should
satisfy the HJ equation.

We define the _Hamilton-Jacobi residual_ and the _Hamilton-Jacobi terminal
condition error_ as follows:

$\displaystyle\text{HJ
residual}:\mathcal{R}_{HJ}(\theta_{U})=\frac{h}{M}\sum_{m=1}^{M}\sum_{k=0}^{K-1}\left|\partial_{t}U(Y_{k}^{(m)},hk;\theta_{U})+\frac{1}{2\lambda}\left|\nabla
U(Y_{k}^{(m)},hk;\theta_{U})\right|^{2}\right|$ (43) $\displaystyle\text{HJ
terminal condition error:
}\mathcal{R}_{T}(\theta_{U})=\frac{1}{M}\sum_{m=1}^{M}\left|\nabla
U(Y_{K}^{(m)},T;\theta_{U})-\nabla\phi^{*}(Y_{K}^{(m)})\right|.$ (44)

These error functions are not used for training, but rather are used as _a
posteriori_ indicators of optimality of the learned generative flow. This is
in constrast to [4], where the Hamilton-Jacobi equation is used as an
additional regularizer during training. Due to Theorem 4.2, there is a unique
(smooth) optimal generative flow and is expected to satisfy the optimality
conditions (20) and (21). Hence any computed generative flow based on
Algorithm 1 can be tested against the corresponding HJ residuals (43) and
(44). In this sense they can also be used as an optimality indicator for
_real-time_ termination of our simulation when the HJ residuals are
sufficiently small. In Section 6, we successfully demonstrate the algorithm on
a MNIST dataset that avoids the use of autoencoders. We show that the
optimality condition are eventually satisfied in the course of our training.
Additional computational costs for calculating (43) and (44) take
$\mathcal{O}(KM)$ and $\mathcal{O}(Md)$ gradient evaluations, respectively,
where $M,K$ are defined in (43) and (44).

###  5.3 Adversarial training eliminates the need for forward-backward
simulation

Continuous Normalizing Flow (CNF) models, such as OT flow [4] requires both
forward and backward simulation of invertible flows during training. These
models learn the flows in the forward normalizing direction

$\inf_{v,\rho}D_{KL}(\rho(\cdot,T),\mathcal{N}_{d})+\lambda\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)dxdt$
(45)

where $\mathcal{N}_{d}$ refers to the $d$-dimensional Gaussian distribution,
and generate samples using the backward flows $\gamma_{t}-\nabla\cdot(\gamma
v)=0$ for $\gamma(\cdot,t)=\rho(\cdot,T-t)$. A main reason for the forward-
backward simulation is because the likelihood needs to be computed by
integrating over the path to estimate the value of the KL divergence. However,
in high-dimensional examples, the data $\pi=\rho_{0}$ often lie on low
dimensional manifolds and immediately two challenges arise.

First, due to the lack of absolute continuity, a KL-ball in (45) fails to
contain an optimizer, leading to unstable training and mathematically to an
undefined terminal condition for the HJ equation $U(\cdot,T)$ if we use
$D_{KL}$ as the terminal condition instead of the Wassersten-1 proximal in
Theorem 4.1. A second mathematical and algorithmic reason for the failure of
CNFs in the case where the target is supported on a lower dimensional
manifold: there is no invertible mapping between the Gaussian and a target
distribution $\pi$ supported on a lower-dimensional distribution, leading to
non-existence of an inverse transport map from $\mathcal{N}_{d}$ to
$\rho_{0}=\pi$. For this reason, autoencoders are necessary for OT flow [4]
and potential flow generators [6] to embed the flows into the appropriate
spaces before applying the generative models.

Unlike CNF-based models such as (45), our model in Algorithm 1, based on
Theorem 4.1, bypasses both such challenges through: (a) _adversarial training_
, which does not require inversion of the flow, and (b) _Wasserstein-1
proximal regularized $f$-divergences_, which do not require absolute
continuity between measures. Therefore our model is able to generate images
without the aid of autoencoders, as we demonstrate next.

##  6 Numerical results

We demonstrate that our proposed generative flows (19) formulated combining
Wasserstein-1 and Wasserstein-2 proximal regularizations facilitate learning
distributions that are supported on low-dimensional manifolds. We learn the
MNIST dataset on $\Omega=[0,1]^{784}$ using the adversarial flow model
described in Algorithm 1, without the use of any autoencoders. To demonstrate
the impact of Wasserstein-1 and Wasserstein-2 proximal regularizations, we
consider the following four test cases:

  1. 1.

Unregularized flow: no running cost penalty for $t\in(0,T)$ and the terminal
cost is the standard dual formulation of $f$-divergences

$\inf_{v}\left\\{\sup_{\phi\in
C(\mathbb{R}^{d})}\left\\{\mathbb{E}_{\rho(\cdot,T)}[\phi]-\mathbb{E}_{\pi}[f^{\star}(\phi)]\right\\}:\frac{dx}{dt}=v(x(t),t),\,x(0)\sim\rho(\cdot,0)\right\\}$
(46)

  2. 2.

Wasserstein-2 proximal of $f$-divergence flow: optimal transport cost for
$t\in(0,T)$ and the terminal cost is the standard dual formulation of
$f$-divergences

$\displaystyle\inf_{v}\left\\{\sup_{\phi\in
C(\mathbb{R}^{d})}\left\\{\mathbb{E}_{\rho(\cdot,T)}[\phi]-\mathbb{E}_{\pi}[f^{\star}(\phi)]\right\\}+\lambda\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)dxdt\right\\}$
(47) $\displaystyle\text{s.t.
}\frac{dx}{dt}=v(x(t),t),\,x(0)\sim\rho(\cdot,0)$

  3. 3.

Wasserstein-1 proximal flow of $f$-divergence flow: no running cost penalty
for $t\in(0,T)$ and the terminal cost is the regularized dual formulation of
$f$-divergences

$\inf_{v}\left\\{\sup_{\phi\in\Gamma_{L}}\left\\{\mathbb{E}_{\rho(\cdot,T)}[\phi]-\mathbb{E}_{\pi}[f^{\star}(\phi)]\right\\}:\frac{dx}{dt}=v(x(t),t),\,x(0)\sim\rho(\cdot,0)\right\\}$
(48)

  4. 4.

Composed Wasserstein-1/Wasserstein-2 proximal flow (W1/W2 proximal flow) :
optimal transport cost for $t\in(0,T)$ and the terminal cost is the
regularized $f$-divergence

$\displaystyle\inf_{v}\left\\{\sup_{\phi\in\Gamma_{L}}\left\\{\mathbb{E}_{\rho(\cdot,T)}[\phi]-\mathbb{E}_{\pi}[f^{\star}(\phi)]\right\\}+\lambda\int_{0}^{T}\int_{\mathbb{R}^{d}}\frac{1}{2}|v(x,t)|^{2}\rho(x,t)dxdt\right\\}$
(49) $\displaystyle\text{s.t.
}\frac{dx}{dt}=v(x(t),t),\,x(0)\sim\rho(\cdot,0)$

For our numerical experiments, we choose the Lipschitz constant to be $L=1$,
the weighting parameters to be $\lambda=0.05$, the terminal time $T=5$, time
step size $h=1.0$, and $f(x)=-\log(x)$ for the reverse KL divergence. In
addition, the discriminator $\phi$ and the potential $U$ share a common neural
network architecture; see Section A.3.

(a) Terminal cost $D_{f}^{\Gamma_{L}}(\rho(\cdot,T)\|\pi)$ (7)

(b) Kinetic energy

(c) Optimality indicator (44)

(d) Optimality indicator (43)

Figure 2: Evaluation of learning objectives (a - b) and optimality indicators
(c - d) from MFG optimality conditions over the course of training. We learn
the MNIST dataset using the four test cases with $T=5.0$, $h=1.0$. We choose
$\lambda=0.05$ and $L=1$ for the weight paramters of Wasserstein-2 and
Wasserstein-1 proximals, respectively. In Figure $\bf{(a)}$, observe that the
terminal cost
$\mathcal{F}(\rho(\cdot,T))=D_{f}^{\Gamma_{L}}(\rho(\cdot,T)\|\pi)$ diverges
(green, red) without the Wasserstein-1 proximal regularization. (b) The
Wasserstein-2 proximal additionally regularizes the flow to have lower kinetic
energy and to be less oscillatory training objectives. Less oscillation is
also related to the uniqueness of the MFG solution in Theorem 4.2 which in
turn is expected to render the algorithms more robust, i.e. in our context not
susceptible to implementation-dependent choices. (c) As inferred from mean-
field game, the learning problem without Wasserstein-1 proximal regularization
lacks a well-defined terminal condition. The exploding optimality indicator
(44) exemplifies this behavior. (d) $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$
proximal generative flow results in lower values of the optimality indicator
(43) compared to those from Wasserstein-1 proximal regularized flow. (c-d)
show that the optimality indicators can inform when an optimal generative flow
has been discovered.

(a) Generated samples from $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ proximal
generative flow (49) with different time step sizes $h=2^{0}$ (left),
$h=2^{-3}$ (center), $h=2^{-6}$ (right)

(b) Generated samples from $\mathcal{W}_{1}$ proximal flow (48) with different
time step sizes $h=2^{0}$ (left), $h=2^{-3}$ (center), $h=2^{-6}$ (right)

Figure 3:  Wasserstein-2 proximal regularization implies discretization
invariance in generative flows. After learning the MNIST dataset
$\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ and $\mathcal{W}_{1}$ proximal
generative flows with $T=5.0$, $h=1.0$, we generated samples by integrating
the learned vector field $\dot{x}(t)=-\frac{1}{\lambda}\nabla U(x(t),t)$ over
time with different time step sizes $h$. In (a), we see that Wasserstein-2
proximal regularization provides almost straight flow trajectories which leads
to generated samples which are almost invariant to time discretization. This
robustness of the continuous time flow generator ensures high fidelity of
generated samples regardless of time discretization. This empirical
observation is justified by the theoretical result in (3) of Theorem 4.1. On
the other hand, we see in (b) that without Wasserstein-2 regularization, the
resulting vector field is more sensitive to varying step sizes as certain
digits may flip to other ones.

Figure 2 demonstrates the indispensability of using the Wasserstein-1 proximal
regularization for stabilizing the training procedure. Observe that the flows
that are not regularized or only rely on Wasserstein-2 regularization _fail_
in the middle of their training procedures as their terminal cost blows up
soon after initialization (see red and green curves in Figure 2(a)). This is
because the terminal cost cannot effectively compare measures $\rho(\cdot,T)$
(which is absolutely continuous with respect to $\rho(\cdot,0)$ and $\pi$. The
Wasserstein-2 proximal does not avoid this issue as it alone cannot make up
for an ill-defined terminal condition.

Moreover, we see in Figure 2(b) that for the flow trained on only the
Wasserstein-1 proximal, the resulting flow has substantially higher kinetic
energy than the the flow that uses both proximals. This implies that the
trajectories are not as regular or potentially strongly nonlinear. Moreover,
it is likely that the learned flow is strongly dependent on the particular
parametrization of the velocity fields.

Our continuous-time flow model stands out for its ability to learn high-
dimensional problems directly in the original space, without specialized
architecture or the use of autoencoders as in normalizng flows. OT-flow [4]
and the potential-flow generator [6] applies autoencoders to get latent
representations of the target in lower dimensional spaces and learn
normalizing flows in the lower dimensional space. Additionally, our flow can
be learned using a discretize-then-optimize (DTO) approach with just a few
time steps $K=5$. We are able to use such a crude step size as the optimal
trajectories are known to be provably linear (Theorem 4.1). This demonstrates
that our problem formulation (19) enables more efficient training compared to
methods with highly engineered architectures such as the 10-layer-stacked CNF
setting used for RNODE [3]. Our approach is similar to the GAN implementation
of the Potential flow generator [6]. But our algorithm differs from it in the
sense that we solve the Wasserstein-2 regularized problem with a Wasserstein-1
regularized terminal cost whose combination provides well-defined optimality
conditions (20), (21) without imposing them as regularization terms in the
objective function.

In Figure 3, we can see that when the two proximals are combined, the
resulting generative flows are _discretization-invariant_ as the generated
images do not change with the step size. This further supporting our claim
that a well-posed generative does not have parametrization-dependent
performance. This behavior is due to Theorem 4.1, which show that the optimal
trajectories are linear, so there is no increased discretization errors with
larger step sizes. On the other hand, when only the $\mathcal{W}_{1}$ proximal
is applied, we see that some images may change with different step sizes,
meaning that the trajectories are incurring extra discretization error with
larger step sizes as they are not linear.

##  7  Discussion and conclusion

In this paper, we introduced $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ proximal
generative flows, which combines Wasserstein-1 and Wasserstein-2 proximal
regularizations for stabilizing the training of continuous-time deterministic
normalizing flows, especially for learning high-dimensional distributions
supported on low-dimensional manifolds. The $\mathcal{W}_{2}$ proximal
regularization enforce flows that minimize kinetic energy, resulting in
straight line paths that are simpler to learn. We showed in Theorem 4.1 that
the optimal trajectories are exactly linear, and we empirically demonstrated
this fact in Figure 3 by showing the trained flow is discretization invariant.

For learning distributions on manifolds, we choose the terminal cost for the
generative flow to be Wasserstein-1 proximal regularization of
$f$-divergences, which can compare mutually singular measures. Typically
targets are supported on an unknown data manifold, so generated models can
easily miss the manifold and be singular with respect to the target
distribution. In addition, Wasserstein-1 proximal regularizations of
$f$-divergences inherit from $f$-divergences the existence of variational
derivatives for any perturbations of measures, including singular ones,
whereas both $f$-divergences and the Wasserstein-1 distance fail in that
respect. This differentiability property suggests that corresponding
minimization problems are smooth and the optimal generative flow can be
robustly discovered by gradient optimization. In the numerical example, we saw
that without a $\mathcal{W}_{1}$ proximal, the training objective would
diverge quickly.

Our mathematical analysis of $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ flows and
their corresponding algorithms, relies on the theory of mean-field game and,
more specifically, on the well-posedness of the MFG optimality conditions,
consisting of a backwards Hamilton-Jacobi equation coupled with a forward
continuity PDE. The potential function for the optimal generative flow will
solve the HJ equation, while the variational derivative of the Wasserstein-1
proximal (9) is exactly the terminal condition of the HJ equation, which is
well-defined for any distributions. We were able to express the optimality
conditions of $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative flows in
Theorem 4.1, meaning that the optimality conditions are at least well-defined.
We then proved in Theorem 4.2, a uniqueness result for the MFG of our
Wasserstein-1 and Wasserstein-2 proximal flow. This showed that the optimal
generative flow is and unique, and therefore the optimization problem is well-
posed, implying that it should be robust to different valid discretizations
and approximations.

We derived an adversarial training algorithm for learning
$\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ proximal generative flows using the
dual formulation of the Wasserstein-1 proximal. Due to this adversarial
structure, no forward-backward simulation of the flow is required, which is in
contrast to likelihood-based training of CNFs. The resulting generative flow
is _more robust_ for learning datasets such as the MNIST dataset without the
use of any pre-trained autoencoders or tuning architectures as is typically
required in CNFs [1, 2, 3, 4, 6].

Finally, we also demonstrated the use of the HJ optimality conditions as a new
real-time criterion for optimality of solutions and termination of the
algorithm. We found them to be more informative for determining when to stop
training the flow rather than imposing them as additional regularizers in the
learning objectives as in some related work [4, 6].

An appealing attribute of $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative
flows is the fact they can be trained with only forward simulations of the
sample trajectories. For future work, we may explore new approaches for
training _stochastic_ normalizing flows, which has been a challenge as
backward SDEs are difficult to simulate efficiently [22, 23]. Moreover, we may
investigate $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative flows and its
connections and applications to distributionally robust optimization (DRO)
[19]. As $\mathcal{W}_{1}\oplus\mathcal{W}_{2}$ generative flows are another
example of an MFG-based generative model, we may rigorously study their
robustness to implementation errors using PDE theory, following the style and
approach of [24].

##  8 Acknowledgement

The authors are partially funded by AFOSR grant FA9550-21-1-0354. H. G., M.K.
and L. R.-B. are partially funded by NSF DMS-2307115. H. G. and M. K. are
partially supported by NSF TRIPODS CISE-1934846.

## References

  * [1] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.  Neural ordinary differential equations.  Advances in neural information processing systems, 31, 2018.
  * [2] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, and David Duvenaud.  Scalable reversible generative models with free-form continuous dynamics.  In International Conference on Learning Representations, 2019.
  * [3] Chris Finlay, Joern-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman.  How to train your neural ODE: the world of Jacobian and kinetic regularization.  In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3154–3164. PMLR, 13–18 Jul 2020\.
  * [4] Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto.  Ot-flow: Fast and accurate continuous normalizing flows via optimal transport.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9223–9232, 2021.
  * [5] Katharina Ott, Prateek Katiyar, Philipp Hennig, and Michael Tiemann.  ResNet after all: Neural ODEs and Their Numerical Solution.  In International Conference on Learning Representations, 2021.
  * [6] Liu Yang and George Em Karniadakis.  Potential flow generator with L2 optimal transport regularity for generative models.  IEEE Transactions on Neural Networks and Learning Systems, 33:528–538, 2019.
  * [7] Jeremiah Birrell, Paul Dupuis, Markos A Katsoulakis, Yannis Pantazis, and Luc Rey-Bellet.  $(f,{\Gamma})$-Divergences: Interpolating between $f$-Divergences and Integral Probability Metrics.  Journal of Machine Learning Research, 23(39):1–70, 2022.
  * [8] Hyemin Gu, Panagiota Birmpa, Yannis Pantazis, Luc Rey-Bellet, and Markos A. Katsoulakis.  Lipschitz-Regularized Gradient Flows and Generative Particle Algorithms for High-Dimensional Scarce Data.  SIAM J.Data Science, to appear, 2024.
  * [9] Benjamin J Zhang and Markos A Katsoulakis.  A mean-field games laboratory for generative modeling.  arXiv preprint arXiv:2304.13534, 2023.
  * [10] Wuchen Li, Siting Liu, and Stanley Osher.  A kernel formula for regularized wasserstein proximal operators.  Research in the Mathematical Sciences, 10(4):43, 2023.
  * [11] Benjamin J Zhang, Siting Liu, Wuchen Li, Markos A Katsoulakis, and Stanley J Osher.  Wasserstein proximal operators describe score-based generative models and resolve memorization.  arXiv preprint arXiv:2402.06162, 2024.
  * [12] Aditya Grover, Manik Dhar, and Stefano Ermon.  Flow-gan: Combining maximum likelihood and adversarial learning in generative models.  In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence and Thirtieth Innovative Applications of Artificial Intelligence Conference and Eighth AAAI Symposium on Educational Advances in Artificial Intelligence, AAAI’18/IAAI’18/EAAI’18. AAAI Press, 2018.
  * [13] Jeremiah Birrell, Yannis Pantazis, Paul Dupuis, Luc Rey-Bellet, and Markos Katsoulakis.  Function-space regularized Rényi divergences.  In The Eleventh International Conference on Learning Representations, 2023.
  * [14] Jeremiah Birrell, Markos A. Katsoulakis, Luc Rey-Bellet, and Wei Zhu.  Structure-preserving GANs.  In International Conference on Machine Learning, 2022.
  * [15] Neal Parikh and Stephen Boyd.  Proximal algorithms.  Found. Trends Optim., 1(3):127–239, Jan 2014.
  * [16] Dávid Terjék.  Moreau-yosida $f$-divergences.  In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10214–10224. PMLR, 18–24 Jul 2021.
  * [17] Wilfrid Gangbo and Michael Westdickenberg.  Optimal transport for the system of isentropic euler equations.  Communications in Partial Differential Equations, 34(9):1041–1073, 2009.
  * [18] Lawrence C Evans.  Partial differential equations, volume 19.  American Mathematical Society, 2022.
  * [19] Hamed Rahimian and Sanjay Mehrotra.  Frameworks and results in distributionally robust optimization.  Open Journal of Mathematical Optimization, 3:1–85, 2022.
  * [20] Jean-Michel Lasry and Pierre-Louis Lions.  Mean field games.  Japanese journal of mathematics, 2(1):229–260, 2007.
  * [21] Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville.  Improved training of Wasserstein GANs.  Advances in neural information processing systems, 30, 2017.
  * [22] Belinda Tzen and Maxim Raginsky.  Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit.  arXiv preprint arXiv:1905.09883, 2019.
  * [23] Liam Hodgkinson, Chris van der Heide, Fred Roosta, and Michael W Mahoney.  Stochastic normalizing flows.  arXiv preprint arXiv:2002.09547, 2020.
  * [24] Nikiforos Mimikos-Stamatopoulos, Benjamin J Zhang, and Markos A Katsoulakis.  Score-based generative models are provably robust: an uncertainty quantification perspective.  arXiv preprint arXiv:2405.15754, 2024.

##  Appendix A Experimental details

###  A.1 Compute resources

Our experiment is computed using personal computer in the environment: Apple
M2 8 cores and Apple M2 24 GB - Metal 3.

###  A.2 Dataset

MNIST is a standard machine learning example where the dataset can be
retrieved from some machine libraries. The original dataset consists of total
60000 samples of $28\times 28$ gray-scale images whose pixel values range from
integers between 0 and 255. We normalize the pixel values and therefore locate
samples in $[0,1]^{28\times 28}$. Each sample represents one of handwritten
digits from 0 to 9. We use randomly chosen 6000 samples (10% of the entire
dataset) as training data.

###  A.3 Neural network architecture

CNN Discriminator
---
$7\times 7$ Conv, $1\times 1$ stride ($1\rightarrow 8$)
ReLU
$2\times 2$ max pool, $2\times 2$ stride
$7\times 7$ Conv, $1\times 1$ stride ($8\rightarrow 8$)
ReLU
$2\times 2$ max pool, $2\times 2$ stride
Flatten with dimension $\ell_{3}$
$W^{4}\in\mathbb{R}^{\ell_{3}\times 256}$, $b^{4}\in\mathbb{R}^{256}$
ReLU
$W^{5}\in\mathbb{R}^{256\times 256}$, $b^{5}\in\mathbb{R}^{256}$
ReLU
$W^{6}\in\mathbb{R}^{256\times 256}$, $b^{6}\in\mathbb{R}^{256}$
ReLU
$W^{7}\in\mathbb{R}^{256\times 1}$, $b^{7}\in\mathbb{R}$
Linear

CNN Potential
---
$7\times 7$ Conv, $1\times 1$ stride ($1\rightarrow 8$)
ReLU
$2\times 2$ max pool, $2\times 2$ stride
$7\times 7$ Conv, $1\times 1$ stride ($8\rightarrow 8$)
ReLU
$2\times 2$ max pool, $2\times 2$ stride
Spatial flatten with dimension $\ell_{3}$
$W^{4}\in\mathbb{R}^{(\ell_{3}+1)\times 512}$, $b^{4}\in\mathbb{R}^{512}$
ReLU
$W^{5}\in\mathbb{R}^{512\times 512}$, $b^{5}\in\mathbb{R}^{512}$
ReLU
$W^{6}\in\mathbb{R}^{512\times 512}$, $b^{6}\in\mathbb{R}^{512}$
ReLU
$W^{7}\in\mathbb{R}^{512\times 1}$, $b^{7}\in\mathbb{R}$
Linear

Table 1: Neural network architectures of discriminator
$\phi:\mathbb{R}^{d}\rightarrow\mathbb{R}$ and potential
$\phi:\mathbb{R}^{d}\times\mathbb{R}\rightarrow\mathbb{R}$. We adapted the
neural network architecture from the code for the potential flow generator GAN
[6] to accommodate high-dimensional image examples.