# GeONet: a neural operator for learning the Wasserstein geodesic

Andrew Gracyk, Xiaohui Chen
Department of Statistics
University of Illinois at Urbana-Champaign
{agracyk2<EMAIL_ADDRESS>

###### Abstract

Optimal transport (OT) offers a versatile framework to compare complex data
distributions in a geometrically meaningful way. Traditional methods for
computing the Wasserstein distance and geodesic between probability measures
require mesh-dependent domain discretization and suffer from the curse-of-
dimensionality. We present _GeONet_ , a mesh-invariant deep neural operator
network that learns the non-linear mapping from the input pair of initial and
terminal distributions to the Wasserstein geodesic connecting the two endpoint
distributions. In the offline training stage, GeONet learns the saddle point
optimality conditions for the dynamic formulation of the OT problem in the
primal and dual spaces that are characterized by a coupled PDE system. The
subsequent inference stage is instantaneous and can be deployed for real-time
predictions in the online learning setting. We demonstrate that GeONet
achieves comparable testing accuracy to the standard OT solvers on a
simulation example and the CIFAR-10 dataset with considerably reduced
inference-stage computational cost by orders of magnitude.

##  1 Introduction

Recent years have seen tremendous progress in statistical and computational
optimal transport (OT) as a lens to explore machine learning problems. One
prominent example is to use the Wasserstein distance to compare data
distributions in a geometrically meaningful way, which has found various
applications, such as in generative models (Arjovsky et al., 2017), domain
adaptation (Courty et al., 2017) and computational geometry (Solomon et al.,
2015). Computing the optimal coupling and the optimal transport map (if it
exists) can be expressed in a fluid dynamics formulation with the minimum
kinetic energy (Benamou and Brenier, 2000). Such dynamical formulation defines
geodesics in the Wasserstein space of probability measures, thus providing
richer information for interpolating between data distributions that can be
used to design efficient sampling methods from high-dimensional distributions
(Finlay et al., 2020). Moreover, the Wasserstein geodesic is also closely
related to the optimal control theory (Chen et al., 2021), which has
applications in robotics and control systems (Krishnan and Martínez, 2018;
Inoue et al., 2021).

Traditional methods for numerically computing the Wasserstein distance and
geodesic require domain discretization that is often mesh-dependent (i.e., on
regular grids or triangulated domains). Classical solvers such as Hungarian
method (Kuhn, 1955), the auction algorithm (Bertsekas and Castanon, 1989), and
transportation simplex (Luenberger and Ye, 2015), suffer from the curse-of-
dimensionality and scale poorly for even moderately mesh-sized problems (Klatt
et al., 2020; Genevay et al., 2016; Benamou and Brenier, 2000). Entropic
regularized OT (Cuturi, 2013) and the Sinkhorn algorithm (Sinkhorn, 1964) have
been shown to efficiently approximate the OT solutions at low computational
cost, handling high-dimensional distributions (Benamou et al., 2015). However,
high accuracy is computationally obstructed with a small regularization
parameter (Altschuler et al., 2017; Dvurechensky et al., 2018). Recently, a
machine learning method to compute the Wasserstein geodesic for a _given_
input pair of probability measures has been considered in (Liu et al., 2021).

Figure 1: A predicted geodesic from our mesh-invariant GeONet at different
spatial resolutions, namely, sampling points (i.e., collocations) are randomly
taken in training and low-resolution inputs can be adapted into high-
resolution geodesics (i.e., super-resolution).

A major challenge of using the OT-based techniques is that one needs to
recompute the Wasserstein distance and geodesic for new input pair of
probability measures. Thus, issues of scalability on large-scale datasets and
suitability in the online learning setting are serious concerns for modern
machine learning, computer graphics, and natural language processing tasks
(Genevay et al., 2016; Solomon et al., 2015; Kusner et al., 2015). This
motivates us to tackle the problem of learning the Wasserstein geodesic from
an _operator learning_ perspective.

There is a recent line of work on learning neural operators for solving
general differential equations or discovering equations from data, including
DeepONet (Lu et al., 2021), Fourier Neural Operators (Li et al., 2020b), and
physics-informed neural networks/operators PINNs (Raissi et al., 2019) and
PINOs (Li et al., 2021). Those methods are mesh-independent, data-driven, and
designed to accommodate specific physical laws governed by certain partial
differential equations (PDEs).

Our contributions. In this paper, we propose a deep neural operator learning
framework _GeONet_ for the Wasserstein geodesic. Our method is based on
learning the optimality conditions in the dynamic formulation of the OT
problem, which is characterized by a coupled PDE system in the primal and dual
spaces. Our main idea is to recast the learning problem of the Wasserstein
geodesic from training data into an operator learning problem for the solution
of the PDEs corresponding to the primal and dual OT dynamics. Our method does
not depend on any pre-specified mesh and it can learn the highly non-linear
Wasserstein geodesic operator from a wide collection of training distributions
that are not necessarily sampled at the same locations in the domain.
Therefore, GeONet is also suitable for zero-shot super-resolution applications
on images, i.e., it is trained on lower resolution and predicts at higher
resolution without seeing any higher resolution data (Shocher et al., 2018).
See Figure 1 for an example of the predicted higher-resolution Wasserstein
geodesic connecting two lower-resolution Gaussian mixture distributions by
GeONet.

Surprisingly, the training of our GeONet does not require the true geodesic
data for connecting the two endpoint distributions. Instead, it only requires
the training data as boundary pairs of initial and terminal distributions. The
reason that GeONet needs much less input data is because its training process
is implicitly informed by the OT dynamics such that the continuity equation in
the primal space and Hamilton-Jacobi equation in the dual space must be
simultaneously satisfied to ensure zero duality gap. Since the geodesic data
are typically difficult to obtain without resorting to some traditional
numerical solvers, the _amortized inference_ nature of GeONet, where inference
on related training pairs can be reused (Gershman and Goodman, 2014), has
substantial computational advantage over standard computational OT methods and
machine learning methods for computing the geodesic designed for single input
pair of distributions (Peyré and Cuturi, 2019; Liu et al., 2021).

Table 1: Comparison of our method GeONet with other existing neural operators and networks for learning dynamics from data. PINN can be found in (Raissi et al., 2019) and the machine learning based minimax method can be found in (Liu et al., 2021). Method characteristic | Classical NNs | PINNs | Minimax | GeONet (Ours)
---|---|---|---|---
no need for retraining for new input |  |  |  |
satisfies the associated PDEs |  |  |  |
does not require known geodesic data |  |  |  |
mesh independence |  |  |  |

Once GeONet training is complete, the inference stage for predicting the
geodesic connecting new initial and terminal data distributions requires only
a forward pass of the network and thus it can be performed in real time. In
contrast, standard OT methods re-compute the Wasserstein distance and geodesic
for each new input distribution pair. This is an appealing feature of
amortized inference to use a pre-trained GeONet for fast geodesic computation
or fine-tuning on a large number of future data distributions. Detailed
comparison between our proposed method GeONet with other existing neural
operators and networks for learning dynamics from data can be found in Table
1.

The rest of the paper is organized as follows. In Section 2, we review some
background on static and dynamic formulations for the OT problem, as well as
on the neural operator learning framework. In Section 3, we present the
proposed GeONet method. In Section 4, we report the numeric performance on a
synthetic experiment and a real image dataset.

##  2 Background

###  2.1 Optimal transport problem: static and dynamic formulations

The optimal mass transportation problem, first considered by the French
engineer Gaspard Monge, is to find an optimal map $T^{*}$ for transporting a
source distribution $\mu_{0}$ to a target distribution $\mu_{1}$ that
minimizes some cost function
$c:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$:

$\min_{T:\mathbb{R}^{d}\to\mathbb{R}^{d}}\left\\{\int_{\mathbb{R}^{d}}c(x,T(x))\;\mathrm{d}\mu_{0}(x):T_{\sharp}\mu_{0}=\mu_{1}\right\\},$
(1)

where $T_{\sharp}\mu$ denotes the pushforward measure defined by
$(T_{\sharp}\mu)(B)=\mu(T^{-1}(B))$ for measurable subset
$B\subset\mathbb{R}^{d}$. In this paper, we focus on the quadratic cost
$c(x,y)={1\over 2}\|x-y\|_{2}^{2}$. The Monge problem (1) induces a metric,
known as the _Wasserstein distance_ , on the space
$\mathcal{P}_{2}(\mathbb{R}^{d})$ of probability measures on $\mathbb{R}^{d}$
with finite second moments. In particular, the 2-Wasserstein distance can be
expressed in the relaxed Kantorovich form:

$W_{2}^{2}(\mu_{0},\mu_{1}):=\min_{\gamma\in\Gamma(\mu_{0},\mu_{1})}\left\\{\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|x-y\|_{2}^{2}\;\mathrm{d}\gamma(x,y)\right\\},$
(2)

where minimization over $\gamma$ runs over all possible couplings
$\Gamma(\mu_{0},\mu_{1})$ with marginal distributions $\mu_{0}$ and $\mu_{1}$.
Problem (2) has the dual form (cf. Villani (2003))

${1\over 2}W_{2}^{2}(\mu_{0},\mu_{1})=\sup_{\varphi\in
L^{1}(\mu_{0}),\;\psi\in
L^{1}(\mu_{1})}\left\\{\int_{\mathbb{R}^{d}}\varphi\;\mathrm{d}\mu_{0}+\int_{\mathbb{R}^{d}}\psi\;\mathrm{d}\mu_{1}:\varphi(x)+\psi(y)\leqslant{\|x-y\|_{2}^{2}\over
2}\right\\}.$ (3)

Problems (1) and (2) are both referred as the _static OT_ problem, which has
close connection to fluid dynamics. Specifically, the Benamou-Brenier dynamic
formulation (Benamou and Brenier, 2000) expresses the Wasserstein distance as
a minimal kinetic energy flow problem:

$\displaystyle\begin{gathered}{1\over
2}W_{2}^{2}(\mu_{0},\mu_{1})=\min_{(\mu,\mathbf{v})}\int_{0}^{1}\int_{\mathbb{R}^{d}}{1\over
2}||\mathbf{v}(x,t)||_{2}^{2}\ \mu(x,t)\ \mathrm{d}x\ \mathrm{d}t\\\
\qquad\qquad\qquad\qquad\mbox{subject to}\ \
\partial_{t}\mu+\mbox{div}(\mu\mathbf{v})=0,\ \ \mu(\cdot,0)=\mu_{0},\ \
\mu(\cdot,1)=\mu_{1},\end{gathered}$ (6)

where $\mu_{t}:=\mu(\cdot,t)$ is the probability density flow at time $t$
satisfying the continuity equation (CE) constraint
$\partial_{t}\mu+\mbox{div}(\mu\mathbf{v})=0$ that ensures the conservation of
unit mass along the flow $(\mu_{t})_{t\in[0,1]}$. To solve (6), we apply the
Lagrange multiplier method to find the saddle point in the primal and dual
variables. In particular, for any flow $\mu_{t}$ initializing from $\mu_{0}$
and terminating at $\mu_{1}$, the Lagrangian function for (6) can be written
as

$\begin{gathered}\mathcal{L}(\mu,\mathbf{v},u)=\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[{1\over
2}\|\mathbf{v}\|_{2}^{2}\mu+\left(\partial_{t}\mu+\mbox{div}(\mu\mathbf{v})\right)u\right]\;\mathrm{d}x\;\mathrm{d}t,\end{gathered}$
(7)

where $u:=u(x,t)$ is the dual variable for the continuity equation. Using
integration-by-parts under suitable decay condition for $\|x\|_{2}\to\infty$,
we find that the optimal dual variable $u^{*}$ satisfies the Hamilton-Jacobi
(HJ) equation for the dynamic OT problem

$\partial_{t}u+{1\over 2}\|\nabla u\|_{2}^{2}=0,$ (8)

and the optimal velocity vector field is given by $\mathbf{v}^{*}(x,t)=\nabla
u^{*}(x,t)$. Hence, we obtained that the Karush–Kuhn–Tucker (KKT) optimality
conditions for (7) are solution $(\mu^{*},u^{*})$ to the following system of
partial differential equations (PDEs):

$\left\\{\begin{gathered}\partial_{t}\mu+\mbox{div}(\mu\nabla u)=0,\ \
\partial_{t}u+{1\over 2}\|\nabla u\|_{2}^{2}=0,\\\ \mu(\cdot,0)=\mu_{0},\ \
\mu(\cdot,1)=\mu_{1}.\end{gathered}\right.$ (9)

In addition, solution to the Hamilton-Jacobi equation (8) can be viewed as an
interpolation $u(x,t)$ of the Kantorovich potential between the initial and
terminal distributions in the sense that $u^{*}(x,1)=\psi^{*}(x)$ and
$u^{*}(x,0)=-\varphi^{*}(x)$ (both up to some additive constants), where
$\psi^{*}$ and $\varphi^{*}$ are the optimal Kantorovich potentials for
solving the static dual OT problem (3). A detailed derivation of the primal-
dual optimality conditions for the dynamical OT formulation is provided in
Appendix A.

###  2.2 Learning neural operators

A neural operator generalizes a neural network that learns a mapping
$\Gamma^{\dagger}:{\cal A}\to{\cal U}$ between infinite-dimensional function
spaces $\cal A$ and $\cal U$ (Kovachki et al., 2021; Li et al., 2020a).
Typically, $\cal A$ and $\cal U$ contain functions defined over a space-time
domain $\Omega\times[0,T]$, where $\Omega$ is taken as a subset of
$\mathbb{R}^{d}$, and the mapping of interest $\Gamma^{\dagger}$ is implicitly
defined through certain differential operator. For example, the physics
informed neural network (PINN) (Raissi et al., 2019) aims to use a neural
network to find a solution to the PDE

$\partial_{t}u+\mathcal{D}[u]=0,$ (10)

given the boundary data $u(\cdot,0)=u_{0}$ and $u(\cdot,T)=u_{T}$, where
$\mathcal{D}:={\cal D}(a)$ denotes a non-linear differential operator in space
that may depend on the input function $a\in{\cal A}$. Different from the
classical neural network learning paradigm that is purely data-driven, a PINN
has less input data (i.e., some randomly sampled data points from the solution
$u=\Gamma^{\dagger}(a)$ and the boundary conditions) since the solution
operator $\Gamma^{\dagger}$ has to obeys the induced physical laws governed by
(10). Even though the PINN is mesh-independent, it only learns the solution
for a single instance of the input function $a$ in the PDE (10). In order to
learn the dynamical behavior of the inverse problem $\Gamma^{\dagger}:{\cal
A}\to{\cal U}$ for an entire family of $\cal A$, we consider the operator
learning perspective.

The idea of using neural networks to approximate any non-linear continuous
operator stems from the universal approximation theorem for operators (Chen
and Chen, 1995; Lu et al., 2021). In particular, we construct a parametric map
by a neural network $\Gamma:\mathcal{A}\times\Theta\rightarrow\mathcal{U}$ for
a finite-dimensional parameter space $\Theta$ to approximate the true solution
operator $\Gamma^{{\dagger}}$. In this paper, we adopt the _DeepONets_
architecture to model $\Gamma$ (Lu et al., 2021). We refer the readers to
Appendix E for some basics of DeepONet and its enhanced versions. Then, the
neural operator learning problem for finding the optimal $\theta^{*}\in\Theta$
can be done in the classical risk minimization framework via

$\displaystyle\begin{gathered}\theta^{*}=\operatorname{argmin}_{\theta\in\Theta}\
\ \
\operatorname{\mathds{E}}_{(a,u_{0},u_{T})\sim\mu}\Big{[}\big{\|}(\partial_{t}+\mathcal{D})\Gamma(a,\theta)\big{\|}_{L^{2}(\Omega\times(0,T))}^{2}\qquad\qquad\qquad\\\
\qquad\qquad\qquad\qquad\qquad+\lambda_{0}\big{\|}\Gamma(a,\theta)(\cdot,0)-u_{0}\big{\|}_{L^{2}(\Omega)}^{2}+\lambda_{T}\big{\|}\Gamma(a,\theta)(\cdot,T)-u_{T}\big{\|}_{L^{2}(\Omega)}^{2}\Big{]},\end{gathered}$
(13)

where the input data $(a,u_{0},u_{T})$ are sampled from some joint
distribution $\mu$. In (13), we minimize the PDE residual corresponding to
$\partial_{t}u+\mathcal{D}[u]=0$ while constraining the network by imposing
boundary conditions. The loss function has weights
$\lambda_{0},\lambda_{T}>0$. Given a finite sample
$\\{(a^{(i)},u_{0}^{(i)},u_{T}^{(i)})\\}_{i=1}^{n}$, and data points randomly
(typically uniformly) sampled in the space-time domain $\Omega\times(0,T)$, we
may minimize the empirical loss analog of (13) by replacing
$\|\cdot\|_{L^{2}(\Omega\times(0,T))}$ with the discrete $L^{2}$ norm over
domain $\Omega\times(0,T)$. Computation of the exact differential operators
$\partial_{t}$ and $\cal D$ can be conveniently exploited via automatic
differentiation in standard deep learning packages.

Neural operators offer advantages over more traditional neural networks that
typically depend on some fixed mesh discretization into a finite-dimensional
space and require the retraining for each different input $(a,u_{0},u_{T})$ to
the dynamical system (10). In contrast, neural operators performs amortized
inference to learn the dynamics for a family of inputs and thus no retraining
is needed for the subsequent inference stage. In addition, neural operators
are flexible enough to transmute context from the original problem into a more
targeted task, producing a discretization-invariant approach that often
achieves more desirable results such as in the zero-shot super-resolution
applications. See more details in Section 4.3.

##  3 Our method

Figure 2: Architecture of GeONet, containing six neural networks to solve the
continuity and Hamilton-Jacobi (HJ) equations, three for each. We minimize the
total loss, and the continuity solution yields the geodesic. GeONet branches
and trunks output vectors of dimension $p$, in which we perform multiplication
among neural network elements to produce the continuity and HJ solutions.

We present GeONet, a geodesic operator network for learning the Wasserstein
geodesic $\\{\mu_{t}\\}_{t\in[0,1]}$ connecting $\mu_{0}$ to $\mu_{1}$ from
the distance $W_{2}(\mu_{0},\mu_{1})$. Let $\Omega\subset\mathbb{R}^{d}$ be
the spatial domain where the probability measures $\mu_{t}$’s are supported.
For probability measures $\mu_{0},\mu_{1}\in{\cal P}_{2}(\Omega)$, it is well-
known that the constant-speed geodesic $\\{\mu_{t}\\}_{t\in[0,1]}$ between
$\mu_{0}$ and $\mu_{1}$ is an absolutely continuous curve in the metric space
$({\cal P}_{2}(\Omega),W_{2})$, which we denote as $\text{AC}({\cal
P}_{2}(\Omega))$. $\mu_{t}$ solves the kinetic energy minimization problem in
(6) (Santambrogio, 2015). Some basic facts on the metric geometry structure of
the Wasserstein geodesic and its relation to the fluid dynamic formulation are
reviewed and discussed in Appendix B. In this work, our goal is to learn a
non-linear operator

$\displaystyle\Gamma^{\dagger}:{\cal P}_{2}(\Omega)\times{\cal
P}_{2}(\Omega)\to$ $\displaystyle\ \text{AC}({\cal P}_{2}(\Omega)),$ (14)
$\displaystyle(\mu_{0},\mu_{1})\mapsto$ $\displaystyle\
\\{\mu_{t}\\}_{t\in[0,1]},$ (15)

based on a training dataset
$\\{(\mu_{0}^{(1)},\mu_{1}^{(1)}),\ldots,(\mu_{0}^{(n)},\mu_{1}^{(n)})\\}$.
The core idea of GeONet is to learn the KKT optimality condition (9) for the
Benamou-Brenier problem. Since (9) is derived to ensure the zero duality gap
between the primal and dual dynamic OT problems, solving the Wasserstein
geodesic requires us to introduce two sets of neural networks that train the
coupled PDEs simultaneously. Specifically, we model the operator learning
problem as an enhanced version of the unstacked DeepONet architecture (Lu et
al., 2021; Tan and Chen, 2022) by jointly training three primal networks in
(16) and three dual networks in (17) as follows:

$\displaystyle\mathcal{C}(x,t,\mu_{0},\mu_{1},\phi)$
$\displaystyle=\sum_{k=1}^{p}\mathcal{B}_{k}^{0,\text{cty}}(\mu_{0},\theta^{0,\text{cty}})\cdot\mathcal{B}_{k}^{1,\text{cty}}(\mu_{1},\theta^{1,\text{cty}})\cdot\mathcal{T}_{k}^{\text{cty}}(x,t,\xi^{\text{cty}}),$
(16) $\displaystyle\mathcal{H}(x,t,\mu_{0},\mu_{1},\psi)$
$\displaystyle=\sum_{k=1}^{p}\mathcal{B}_{k}^{0,\text{HJ}}(\mu_{0},\theta^{0,\text{HJ}})\cdot\mathcal{B}_{k}^{1,\text{HJ}}(\mu_{1},\theta^{1,\text{HJ}})\cdot\mathcal{T}_{k}^{\text{HJ}}(x,t,\xi^{\text{HJ}}),$
(17)

where
$\mathcal{B}^{j,\text{cty}}(\mu_{0}(x_{1}),\dots,\mu_{0}(x_{m}),\theta^{j,\text{cty}}):\mathbb{R}^{m}\times\Theta\rightarrow\mathbb{R}^{p}$
and
$\mathcal{B}^{j,\text{HJ}}(\mu_{0}(x_{1}),\dots,\mu_{0}(x_{m}),\theta^{j,\text{HJ}}):\mathbb{R}^{m}\times\Theta\rightarrow\mathbb{R}^{p}$
are _branch_ neural networks taking $m$-discretized input of initial and
terminal density values at $j=0$ and $j=1$ respectively, and
$\mathcal{T}^{\text{cty}}(x,t,\xi^{\text{cty}}):\mathbb{R}^{d}\times[0,1]\times\Xi\rightarrow\mathbb{R}^{p}$
and
$\mathcal{T}^{\text{HJ}}(x,t,\xi^{\text{HJ}}):\mathbb{R}^{d}\times[0,1]\times\Xi\rightarrow\mathbb{R}^{p}$
are _trunk_ neural networks taking spatial and temporal inputs (cf. Appendix E
for more details on DeepONet models). Here $\Theta$ and $\Xi$ are finite-
dimensional parameter spaces, and $p$ is the output dimension of the branch
and truck networks. Denote parameter concatenations
$\phi:=(\theta^{0,\text{cty}},\theta^{1,\text{cty}},\xi^{\text{cty}})$ and
$\psi:=(\theta^{0,\text{HJ}},\theta^{1,\text{HJ}},\xi^{\text{HJ}})$. Then the
primal operator network
$\mathcal{C}_{\phi}(x,t,\mu_{0},\mu_{1}):=\mathcal{C}(x,t,\mu_{0},\mu_{1},\phi)$
for $\phi\in\Theta\times\Theta\times\Xi$ acts as an approximate solution to
the continuity equation, hence the true geodesic
$\Gamma^{{\dagger}}(x,t,\mu_{0}(x),\mu_{1}(x)):=\mu_{t}(x)=\mu(x,t)$, while
the dual operator network $\mathcal{H}_{\psi}(x,t,\mu_{0},\mu_{1})$ for
$\psi\in\Theta\times\Theta\times\Xi$ corresponds to that of the associated
Hamilton-Jacobi equation. The architecture of GeONet is shown in Figure 2.

To train the GeONet defined in (16) and (17), we minimize the empirical loss
function corresponding to the system of primal-dual PDEs and boundary
residuals in (9) over the parameter spaces $\Theta$ and $\Xi$:

$\phi^{*},\psi^{*}\ \ \ =\ \ \
\operatorname{argmin}_{\phi,\psi\in\Theta\times\Theta\times\Xi}\ \ \
\mathcal{L}_{\text{cty}}+\mathcal{L}_{\text{HJ}}+\mathcal{L}_{\text{BC}},$
(18)

where

$\displaystyle\mathcal{L}_{\text{cty},i}\ \ \ $ $\displaystyle=\ \ \
\frac{\alpha_{1}}{N}\ \Big{\|}\ \ \ \frac{\partial}{\partial
t}\mathcal{C}_{\phi,i}+\mbox{div}(\mathcal{C}_{\phi,i}\nabla\mathcal{H}_{\psi,i})\
\ \ \Big{\|}_{L^{2}(\Omega\times(0,1))}^{2},$ (19)
$\displaystyle\mathcal{L}_{\text{HJ},i}\ \ \ $ $\displaystyle=\ \ \
\frac{\alpha_{2}}{N}\ \Big{\|}\ \ \ \frac{\partial}{\partial
t}\mathcal{H}_{\psi,i}+\frac{1}{2}||\nabla\mathcal{H}_{\psi,i}||_{2}^{2}\ \ \
\Big{\|}_{L^{2}(\Omega\times(0,1))}^{2},$ (20)
$\displaystyle\mathcal{L}_{\text{BC},i}\ \ \ $ $\displaystyle=\ \ \
\frac{\beta_{0}}{N}\ \Big{\|}\ \mathcal{C}_{\phi,0,i}-\mu_{0}^{(i)}\
\Big{\|}_{L^{2}(\Omega)}^{2}\ \ +\ \ \frac{\beta_{1}}{N}\ \Big{\|}\
\mathcal{C}_{\phi,1,i}-\mu_{1}^{(i)}\ \Big{\|}_{L^{2}(\Omega)}^{2},$ (21)

$\mathcal{L}_{\text{cty}}=\sum_{i=1}^{n}\mathcal{L}_{\text{cty},i}$,
$\mathcal{L}_{HJ}=\sum_{i=1}^{n}\mathcal{L}_{\text{HJ},i}$ and
$\mathcal{L}_{\text{BC}}=\sum_{i=1}^{n}\mathcal{L}_{\text{BC},i}$. Here,
$\mathcal{C}_{\phi,i}(x,t):=\mathcal{C}_{\phi}(x,t,\mu_{0}^{(i)}(x),\mu_{1}^{(i)}(x))$
and
$\mathcal{C}_{\phi,t,i}(x):=\mathcal{C}_{\phi}(x,t,\mu_{0}^{(i)}(x),\mu_{1}^{(i)}(x))$
denote the evaluation of neural network $\mathcal{C}_{\phi}$ over the $i$-th
distribution of the $n$ training data at space location $x$ and time point
$t$. The same notation applies for the Hamilton-Jacobi neural networks.
$\mathcal{L}_{\text{cty}}$ is the loss component in which the continuity
equation is satisfied, and $\mathcal{L}_{\text{HJ}}$ is the Hamilton-Jacobi
equation loss component. The boundary conditions are incorporated in the
$\mathcal{L}_{\text{BC}}$ term, and
$\alpha_{1},\alpha_{2},\beta_{1},\beta_{2}$ are scaling weights for the
strength to impose the physics-informed loss. Automatic differentiation of our
GeONet involves differentiating the coupled DeepONet architecture (cf. Figure
2) in order to compute the physics-informed loss terms, which we discuss in
Appendix F.

In practice, we only have access to point clouds of the training distribution
data $(\mu_{0},\mu_{1})$ and the Wasserstein geodesic is output on a
discretized space-time domain. Following (Raissi et al., 2019), we utilize a
_collocation_ procedure as follows. We sample pairs $(x,t)$ randomly and
uniformly within the bounded domain $\Omega\times[0,1]$, i.e.,
$U(\Omega)\times U([0,1])$. These pairs are resampled during each epoch in our
method. Then, we evaluate the continuity and Hamilton-Jacobi residuals
displayed in (19) and (20) at such sampled values, in which the loss is
subsequently minimized.

In GeONet, we implement the fully connected artificial neural networks (ANNs)
for the branch and trunk networks. In our collocation procedure, we take
equispaced locations $x_{1},\ldots,x_{m}$ within $\Omega$, a bounded domain in
$\mathbb{R}^{d}$, typically a hypercube, for branch input. We additionally
take $M$ spatial and temporal locations $(x,t)$ for the trunk, but these are
evaluated randomly and uniformly among the hypercube, different for each
density pair $(\mu_{0},\mu_{1})$. The total number of training data is $N=nM$.
GeONet training pseudo-code is in Algorithm 1.

Algorithm 1 End-to-end training of GeONet

Input: data pairs
$(\mu_{0}^{(1)}\mu_{1}^{(1)}),\ldots,(\mu_{0}^{(n)},\mu_{1}^{(n)})$;
discretization size $N$; initialization of the neural network parameters
$\phi,\psi\in\Theta\times\Theta\times\Xi$; weight parameters
$\alpha_{1},\alpha_{2},\beta_{0},\beta_{1}$.

1:while $\mathcal{L}_{\text{total}}$ has not converged do

2: Independently draw $N$ sample points
$\\{(x_{\ell}^{i},t_{\ell}^{i})\\}_{\ell=1}^{N}$ from $U(\Omega)\times U(0,1)$

3: Compute
$\Phi_{i}=\partial_{t}\mathcal{C}_{\phi,i}+\mbox{div}(\mathcal{C}_{\phi,i}\nabla\mathcal{H}_{\psi,i})$.
$\triangleright$ continuity residual

4: Compute $\Psi_{i}=\partial_{t}\mathcal{H}_{\psi,i}+{1\over
2}\|\mathcal{H}_{\psi,i}\|_{2}^{2}$. $\triangleright$ Hamilton-Jacobi residual

5: Compute $B_{0,i}=\mathcal{C}_{\phi,0,i}-\mu_{0}^{(i)}(x_{\ell}^{i}),\ \
B_{1,i}=\mathcal{C}_{\phi,1,i}-\mu_{1}^{(i)}(x_{\ell}^{i})$. $\triangleright$
boundary residual

6: Compute
$\begin{gathered}\mathcal{L}_{\text{cty}}=\frac{\alpha_{1}}{N}\sum_{i=1}^{n}\|\Phi_{i}\|_{L^{2}(\Omega\times(0,1))}^{2},\
\ \ \
\mathcal{L}_{\text{HJ}}=\frac{\alpha_{2}}{N}\sum_{i=1}^{n}\|\Psi_{i}\|_{L^{2}(\Omega\times(0,1))}^{2},\\\
\mathcal{L}_{\text{BC}}=\frac{1}{N}\sum_{i=1}^{n}\beta_{0}\|B_{0,i}\|_{L^{2}(\Omega)}^{2}+\beta_{1}\|B_{1,i}\|_{L^{2}(\Omega)}^{2},\end{gathered}$
where $\|f\|_{L^{2}(\Omega\times(0,1))}^{2}:=\sum_{k=1}^{M}f^{2}(x_{k},t_{k})$
and $\|f\|_{L^{2}(\Omega)}^{2}:=\sum_{k=1}^{M}f^{2}(x_{k})$.

7: Compute
$\mathcal{L}_{\text{total}}(\phi,\psi)=\mathcal{L}_{\text{cty}}+\mathcal{L}_{\text{HJ}}+\mathcal{L}_{\text{BC}}$.

8: Minimize $\mathcal{L}_{\text{total}}(\phi,\psi)$ to update $\phi$ and
$\psi$. $\triangleright$ minimize the loss function

9:end while

Entropic regularization. Our GeONet is compatible with entropic
regularization, which is related to the Schrödinger bridge problem and
stochastic control (Chen et al., 2016). In Appendix C, we propose the
_entropic-regularized GeONet_ (ER-GeONet), which learns a similar system of
KKT conditions for the optimization as in (9). In the zero-noise limit as the
entropic regularization parameter $\varepsilon\downarrow 0$, the solution of
the optimal entropic interpolating flow converges to solution of the Benamou-
Brenier problem (6) in the sense of the method of vanishing viscosity (Mikami,
2004; Evans, 2010).

Gradient enhancement. In practice, we fortify the base method by adding extra
residual terms of the differentiated PDEs to our loss function of GeONet. Such
gradient enhancement technique has been used to strengthen PINNs (Yu et al.,
2022), which improves efficiency as fewer data points are needed to be sampled
from $U(\Omega)\times U(0,1)$, and prediction accuracy as well. We introduce
the gradient-enhanced version of GeONet with more details in Appendix D.

##  4 Numerical experiments

###  4.1 Gaussian mixture distributions

Since finite mixture distributions are powerful universal approximator for
continuous probability density functions (Nguyen et al., 2020), we first
deploy GeONet on Gaussian mixture distributions over two-dimensional domains.
We learn the Wasserstein geodesic mapping between initial and terminal
distributions of the form
$\mu_{j}(x)=\sum_{i=1}^{k_{j}}\pi_{i}\mathcal{N}(x|u_{i},\Sigma_{i})$ subject
to $\sum_{i=1}^{k_{j}}\pi_{i}=1$, where $j\in\\{0,1\\}$ corresponds to initial
and terminal distributions $\mu_{0},\mu_{1}$, and $k_{j}$ denotes the number
of components taken in the mixture. Here $u_{i}$ and $\Sigma_{i}$ are the mean
vectors and covariance matrices of individual Gaussian components,
respectively.

Figure 3: A geodesic predicted by GeONet and a reference geodesic computed by POT on a test Gaussian mixture distribution pair $(\mu_{0},\mu_{1})$ with $k_{0}=5$ (with stacked means) and $k_{1}=2$. The reference serves as a close approximation to the true geodesic due to entropic regularization. Gaussian means are shown in the references at $t=0,1$.  Table 2: MSE of GeONet on 50 test data of Gaussian mixtures over a $50\times 50$ mesh. All values are multiplied by $10^{-3}$. We report the means and standard deviations of the MSE. Training is done with fixed $k_{0}=5,k_{1}=2$ in all cases with $\pi_{i}$ randomly chosen. Training data at the boundaries has resolution $30\times 30$.  | GeONet error for Gaussian mixtures
---|---
Number of Gaussians | $\bm{t=0}$ | $\bm{t=0.25}$ | $\bm{t=0.5}$ | $\bm{t=0.75}$ | $\bm{t=1}$
Identity $k_{0}=k_{1}=2$ | $0.44\pm 0.44$ | $1.40\pm 0.73$ | $2.10\pm 1.20$ | $1.40\pm 0.70$ | $0.42\pm 0.43$
$k_{0}=3,k_{1}=2$ | $0.39\pm 0.33$ | $2.60\pm 2.80$ | $4.40\pm 5.10$ | $2.90\pm 2.90$ | $0.57\pm 0.61$
$k_{0}=4,k_{1}=2$ | $0.43\pm 0.37$ | $2.90\pm 2.60$ | $5.10\pm 4.70$ | $3.00\pm 2.70$ | $0.42\pm 0.40$
$k_{0}=5,k_{1}=2$ | $0.33\pm 0.22$ | $2.20\pm 1.80$ | $4.10\pm 3.50$ | $2.60\pm 2.20$ | $0.44\pm 0.51$
OOD $k_{0}=2,k_{1}=5$ | $1.20\pm 0.51$ | $9.30\pm 3.60$ | $16.0\pm 7.70$ | $11.0\pm 5.00$ | $11.0\pm 4.80$

Experiment details. In our experiment, domain
$\Omega=[0,5]\times[0,5]\subseteq\mathbb{R}^{2}$ was chosen, which was
discretized into a $30\times 30$ grid for GeONet input, meaning the branch
networks took vector input of $900$ in dimension each. We generate $1,500$
training pairs $(\mu_{0},\mu_{1})$ with $900$ collocations for each pair,
yielding a total of $1.35\times 10^{6}$ training data points. Recall that
GeONet is mesh-invariant, so the $30\times 30$ grids can be adapted to any
higher resolution. In Table 4, mean vectors were $u_{i}\in[1.3,3.7]^{2}$,
variance $\sigma_{0,i}^{2},\sigma_{1,i}^{2}\in[0.4,1.0]$, and covariance
$\sigma_{01,i}\in[-0.4,0.4]$. Additional training details are given in
Appendix H.

We use the mean squared error (MSE) $\int_{\Omega}|\mathcal{C}-\mu|^{2}dx$ as
our error metric to assess the performance, where $\mu$ is a reference
geodesic as proxy of the true geodesic without entropic regularization. This
integral is made discrete by evaluating a Riemann sum along a mesh. The
reference is computed using the Convolutional Wasserstein Barycenter framework
within the POT Python library (Solomon et al., 2015; Flamary et al., 2021).
POT takes a regularization parameter, which "smooths" the solutions. In our
comparison, we choose $\nu=0.003$.

Performance. Our baseline results were collected by probing GeONet on the
identity geodesic by choosing $\mu_{0}=\mu_{1}$ with $k_{0}=k_{1}=2$ in Table
4. The baseline identity geodesic provides a benchmark for comparing and
interpreting the MSEs across different setups. From Table 4, we can draw the
following observations. First, prediction accuracy is the worst at $t=0.5$
(i.e., the barycenter of $\mu_{0}$ and $\mu_{1}$), while the loss boundary
conditions (21) allow greater precision for $t=0,1$, which suggests that lack
of data-enforced conditions along the inner region of the time continuum would
cause greater error. Visualization of a geodesic predicted by GeONet and a
reference geodesic computed by POT is shown in Figure 3. Second, since the
testing pairs $k_{0}=2,3,4$ with fixed $k_{1}=2$ can be viewed as degenerate
cases of the training setting $k_{0}=5,k_{1}=2$, the _in-distribution_ MSEs
are comparable across these setups. Third, MSEs for predicting the trivial
identity geodesic in the intermediate $t=0.25,0.5,0.75$ are uniformly smaller
than other in-distribution setups since the former is an easier task. Finally,
we performed a similar experiment with training and testing on the same
Gaussian mixture distribution family with $k_{0}=k_{1}=5$ and $\pi_{i}=0.2$
for all $i$. Similar conclusions can be drawn, and the MSEs are presented in
Appendix G.

Figure 4: OOD generalization at test time with $k_{0}=2,k_{1}=5$.

Out-of-distribution generalization. To examije the generalization performance
of GeONet under distribution shifts, we consider a setup with
$k_{0}=2,k_{1}=5$ which is out-of-distribution (OOD) at test time. The last
row in Table 4 shows that the MSEs are significantly higher than the in-
distribution settings. Figure 4 shows an example for comparing the barycenter
at $t=0.5$ predicted from GeONet and that computed from POT, where the latter
serves as the reference to ground truth.

###  4.2 A real data application

Our next experiment was upon the CIFAR-10 dataset $32\times 32$ images. As in
Section 4.1, we chose domain $\Omega=[0,5]\times[0,5]$, in which the image RGB
color channels are normalized. GeONet took image input of $32\times 32=1,024$
vectors for each branch network. More training details can be found in
Appendix H. In our experiment, we stack the channel-wise continuity and
Hamilton-Jacobi networks to learn the geodesics for all three channels
simultaneously and combine three geodesics for each normalized RGB color
channel to produce the output images.

Table 3: MSE of GeONet on 50 test pairings of CIFAR-10 data. All values are multiplied by $10^{-3}$. Errors are averaged across all three channels after normalization as densities. | GeONet error on CIFAR-10 data
---|---
Testing setting | $\bm{t=0}$ | $\bm{t=0.25}$ | $\bm{t=0.5}$ | $\bm{t=0.75}$ | $\bm{t=1}$
Identity pairing | $2.0\pm 1.7$ | $3.3\pm 3.5$ | $5.8\pm 5.8$ | $3.9\pm 4.1$ | $3.6\pm 4.1$
Random pairing | $2.1\pm 1.1$ | $4.4\pm 3.8$ | $7.4\pm 6.8$ | $4.7\pm 4.5$ | $4.0\pm 6.0$

Performance. Averaged MSEs of GeONet over the three RGB channels on 50 testing
pairs from CIFAR-10 data are shown in Table 3. We observe that the testing
errors for identity and random pairs are similar. As in the Gaussian mixture
experiment, MSEs are smaller when we move closer to the geodesic boundaries at
$t=0,1$.

Regularization. Classical algorithms such as POT take regularization parameter
$\nu$, which results in a smoothing and approximate of the solutions. This
characteristic is pronounced over a coarse mesh with fluctuating values, such
as values in CIFAR images, which induces larger geodesic inaccuracy. In
contrast, GeONet takes no regularization, which is beneficial towards better
accuracy.

###  4.3 Zero-shot super-resolution application

GeONet is mesh-invariant and suitable for zero-shot super-resolution examples
of adapting low-resolution data into high-resolution geodesics, which includes
initial data at $t=0,1$. The neural operator can be trained only on such low-
resolution information, and the super-resolution application results from the
uniform sampling done within the training and physics-informed loss.
Traditional OT solvers have no ability to do this, as they are confined to the
original mesh. Figure 5 (left) shows examples for two lower-resolution
CIFAR-10 images adapted by GeONet to connect on a geodesic of higher-
resolution, where each row represents the original and higher-resolution
initial and terminal images.

Figure 5: Left: display of zero-shot super-resolution CIFAR-10 examples.
Lower- to higher-resolution geodesics on each row are produced, which includes
the input images at $t=0,1$. Right: comparison of log-runtime in seconds on
Gaussian mixtures. We average over 50 pairs with varying mesh sizes.

###  4.4 Runtime comparison

Inference of GeONet is instantaneous, a feature advantageous for many pairs
and high resolution images. POT is greatly encumbered by a fine mesh. Figure 5
(right) demonstrates an inference-stage runtime comparison of GeONet to the
POT algorithm on Gaussian mixtures considered in Section 4.1. The POT
algorithm was minimally modified, only so that the pertinent distributions
could be taken as input. The linear pattern of Figure 5 (right) on log-scale
suggests the computational improvement of GeONet over POT is of the orders of
magnitude. We remark that CIFAR-10 data in our experiment setup has the
problem size corresponding to $x=10.8$. Training time for both Gaussian
mixtures and CIFAR-10 experiments was approximately $15$ hours on a single
NVIDIA Tesla P100 GPU. Empirically, we found training time to be faster if
coefficients $\alpha_{1},\alpha_{2}$ were scaled similarly to
$\beta_{0},\beta_{1}$ (i.e., more balanced physics and data losses), but some
accuracy is lost. In particular, training time for the Gaussian experiment was
approximately $8$ hours for $\alpha_{1}=\alpha_{2}=\beta_{0}=\beta_{1}$ = 1.
More training details can be found in Appendix H.

## References

  * Altschuler et al. [2017] Jason Altschuler, Jonathan Weed, and Philippe Rigollet.  Near-linear time approximation algorithms for optimal transport via sinkhorn iteration.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , NIPS’17, page 1961–1971, Red Hook, NY, USA, 2017. Curran Associates Inc.  ISBN 9781510860964.
  * Ambrosio et al. [2008] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré.  _Gradient Flows in Metric Spaces and in the Space of Probability Measures_.  Lectures in Mathematics ETH Zürich. Birkhäuser Baseluser Basel, second edition, 2008.
  * Arjovsky et al. [2017] Martin Arjovsky, Soumith Chintala, and Léon Bottou.  Wasserstein generative adversarial networks.  In Doina Precup and Yee Whye Teh, editors, _Proceedings of the 34th International Conference on Machine Learning_ , volume 70 of _Proceedings of Machine Learning Research_ , pages 214–223. PMLR, 06–11 Aug 2017.  URL https://proceedings.mlr.press/v70/arjovsky17a.html.
  * Benamou and Brenier [2000] Jean-David Benamou and Yann Brenier.  A computational fluid mechanics solution to the monge-kantorovich mass transfer problem.  _Numerische Mathematik_ , 84:375–393, 2000.
  * Benamou et al. [2015] Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré.  Iterative bregman projections for regularized transportation problems.  _SIAM Journal on Scientific Computing_ , 37(2):A1111–A1138, 2015.  doi: 10.1137/141000439.  URL https://doi.org/10.1137/141000439.
  * Bertsekas and Castanon [1989] Dimitri P. Bertsekas and David A. Castanon.  The auction algorithm for the transportation problem.  _Annals of Operations Research_ , 20(1):67–96, December 1989.  ISSN 0254-5330.  doi: 10.1007/BF02216923.
  * Brenier [1991] Yann Brenier.  Polar factorization and monotone rearrangement of vector-valued functions.  _Communications on Pure and Applied Mathematics_ , 44(4):375–417, 1991.  doi: https://doi.org/10.1002/cpa.3160440402.  URL https://onlinelibrary.wiley.com/doi/abs/10.1002/cpa.3160440402.
  * Burage et al. [2001] Dmitri Burage, Yuri Burago, and Sergei Ivanov.  _A Course in Metric Geometry_.  Graduate studies in mathematics. American mathematical society, Providence, Rhode Island, 2001.
  * Chen and Chen [1995] Tianping Chen and Hong Chen.  Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems.  _IEEE Transactions on Neural Networks_ , 6(4):911–917, 1995.  doi: 10.1109/72.392253.
  * Chen and Yang [2021] Xiaohui Chen and Yun Yang.  Cutoff for exact recovery of gaussian mixture models.  _IEEE Transactions on Information Theory_ , 67(6):4223–4238, 2021.  doi: 10.1109/TIT.2021.3063155.
  * Chen et al. [2016] Yongxin Chen, Tryphon T. Georgiou, and Michele Pavon.  On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint.  _Journal of Optimization Theory and Applications_ , 169(2):671–691, 2016.  doi: 10.1007/s10957-015-0803-z.  URL https://doi.org/10.1007/s10957-015-0803-z.
  * Chen et al. [2021] Yongxin Chen, Tryphon T. Georgiou, and Michele Pavon.  Optimal transport in systems and control.  _Annual Review of Control, Robotics, and Autonomous Systems_ , 4(1):89–113, 2021.  doi: 10.1146/annurev-control-070220-100858.  URL https://doi.org/10.1146/annurev-control-070220-100858.
  * Courty et al. [2017] Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy.  Optimal transport for domain adaptation.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 39(9):1853–1865, 2017.  doi: 10.1109/TPAMI.2016.2615921.
  * Cuturi [2013] Marco Cuturi.  Sinkhorn distances: Lightspeed computation of optimal transport.  In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger, editors, _Advances in Neural Information Processing Systems_ , volume 26. Curran Associates, Inc., 2013.  URL https://proceedings.neurips.cc/paper/2013/file/af21d0c97db2e27e13572cbf59eb343d-Paper.pdf.
  * Dvurechensky et al. [2018] Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin.  Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn’s algorithm.  In Jennifer Dy and Andreas Krause, editors, _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pages 1367–1376. PMLR, 10–15 Jul 2018.
  * Evans [2010] Lawrence C. Evans.  _Partial Differential Equations_.  American Mathematical Society, Providence, R.I., 2010.  ISBN 9780821849743 0821849743.
  * Finlay et al. [2020] Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam M Oberman.  "how to train your neural ode: the world of jacobian and kinetic regularization".  In _ICML_ , 2020.
  * Flamary et al. [2021] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer.  Pot: Python optimal transport.  _Journal of Machine Learning Research_ , 22(78):1–8, 2021.
  * Genevay et al. [2016] Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach.  Stochastic optimization for large-scale optimal transport.  In _Proceedings of the 30th International Conference on Neural Information Processing Systems_ , NIPS’16, page 3440–3448, Red Hook, NY, USA, 2016. Curran Associates Inc.  ISBN 9781510838819.
  * Gershman and Goodman [2014] Samuel J. Gershman and Noah D. Goodman.  Amortized inference in probabilistic reasoning.  _Cognitive Science_ , 36, 2014.
  * Inoue et al. [2021] Daisuke Inoue, Yuji Ito, and Hiroaki Yoshida.  Optimal transport-based coverage control for swarm robot systems: Generalization of the voronoi tessellation-based method.  In _2021 American Control Conference (ACC)_ , pages 3032–3037, 2021\.  doi: 10.23919/ACC50511.2021.9483194.
  * Klatt et al. [2020] Marcel Klatt, Carla Tameling, and Axel Munk.  Empirical regularized optimal transport: Statistical theory and applications.  _SIAM Journal on Mathematics of Data Science_ , 2(2):419–443, 2020.  doi: 10.1137/19M1278788.
  * Kovachki et al. [2021] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Neural operator: Learning maps between function spaces, 2021.  URL https://arxiv.org/abs/2108.08481.
  * Krishnan and Martínez [2018] Vishaal Krishnan and Sonia Martínez.  Distributed optimal transport for the deployment of swarms.  In _2018 IEEE Conference on Decision and Control (CDC)_ , pages 4583–4588, 2018.  doi: 10.1109/CDC.2018.8619816.
  * Kuhn [1955] H. W. Kuhn.  The hungarian method for the assignment problem.  _Naval Research Logistics Quarterly_ , 2(1-2):83–97, 1955.  doi: https://doi.org/10.1002/nav.3800020109.  URL https://onlinelibrary.wiley.com/doi/abs/10.1002/nav.3800020109.
  * Kusner et al. [2015] Matt J. Kusner, Yu Sun, Nicholas I. Kolkin, and Kilian Q. Weinberger.  From word embeddings to document distances.  In _Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37_ , ICML’15, page 957–966. JMLR.org, 2015.
  * Li et al. [2020a] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Neural operator: Graph kernel network for partial differential equations, 2020a.  URL https://arxiv.org/abs/2003.03485.
  * Li et al. [2020b] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Fourier neural operator for parametric partial differential equations, 2020b.  URL https://arxiv.org/abs/2010.08895.
  * Li et al. [2021] Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar.  Physics-informed neural operator for learning partial differential equations, 2021.  URL https://arxiv.org/abs/2111.03794.
  * Liu et al. [2021] Shu Liu, Shaojun Ma, Yongxin Chen, Hongyuan Zha, and Haomin Zhou.  Learning high dimensional wasserstein geodesics, 2021.  URL https://arxiv.org/abs/2102.02992.
  * Lu et al. [2021] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis.  Learning nonlinear operators via deeponet based on the universal approximation theorem of operators.  _Nat Mach Intell_ , 3:218–229, 2021.
  * Luenberger and Ye [2015] David G. Luenberger and Yinyu Ye.  _Linear and Nonlinear Programming_.  Springer Publishing Company, Incorporated, 2015.  ISBN 3319188410.
  * McCann [1997] Robert J. McCann.  A convexity principle for interacting gases.  _Advances in Mathematics_ , 128(1):153–179, 1997\.  URL https://www.sciencedirect.com/science/article/pii/S0001870897916340.
  * Mikami [2004] Toshio Mikami.  Monge’s problem with a quadratic cost by the zero-noise limit of h-path processes.  _Probability Theory and Related Fields_ , 129(2):245–260, 2004.  doi: 10.1007/s00440-004-0340-4.  URL https://doi.org/10.1007/s00440-004-0340-4.
  * Nguyen et al. [2020] T. Tin Nguyen, Hien D. Nguyen, Faicel Chamroukhi, and Geoffrey J. McLachlan.  Approximation by finite mixtures of continuous density functions that vanish at infinity.  _Cogent Mathematics & Statistics_, 7(1):1750861, 2020.  doi: 10.1080/25742558.2020.1750861.  URL https://doi.org/10.1080/25742558.2020.1750861.
  * Peyré and Cuturi [2019] Gabriel Peyré and Marco Cuturi.  Computational optimal transport.  _Foundations and Trends in Machine Learning_ , 11(5-6):355–607, 2019.
  * Raissi et al. [2019] M. Raissi, P. Perdikaris, and G.E. Karniadakis.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.  _Journal of Computational Physics_ , 378:686–707, 2019\.  ISSN 0021-9991.  doi: https://doi.org/10.1016/j.jcp.2018.10.045.  URL https://www.sciencedirect.com/science/article/pii/S0021999118307125.
  * Santambrogio [2015] Filippo Santambrogio.  _Optimal Transport for Applied Mathematicians. Calculus of Variations, PDEs and Modeling_.  Birkhäuser Basel, 1 edition, 2015.  URL https://www.math.u-psud.fr/~filippo/OTAM-cvgmt.pdf.
  * Shocher et al. [2018] Assaf Shocher, Nadav Cohen Cohen, and Michal Irani.  “zero-shot" super-resolution using deep internal learning.  In _The IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , June 2018.
  * Sinkhorn [1964] Richard Sinkhorn.  A Relationship Between Arbitrary Positive Matrices and Doubly Stochastic Matrices.  _The Annals of Mathematical Statistics_ , 35(2):876 – 879, 1964.  doi: 10.1214/aoms/1177703591.  URL https://doi.org/10.1214/aoms/1177703591.
  * Solomon et al. [2015] Justin Solomon, Fernando de Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas.  Convolutional wasserstein distances: Efficient optimal transportation on geometric domains.  _ACM Trans. Graph._ , 34(4), jul 2015.  ISSN 0730-0301.  doi: 10.1145/2766963.  URL https://doi.org/10.1145/2766963.
  * Tan and Chen [2022] Lesley Tan and Liang Chen.  Enhanced deeponet for modeling partial differential operators considering multiple input functions, 2022.  URL https://arxiv.org/abs/2202.08942.
  * Villani [2003] Cédric Villani.  _Topics in optimal transportation_.  Graduate studies in mathematics. American mathematical society, Providence, Rhode Island, 2003.  ISBN 0-8218-3312-X.
  * Wang et al. [2021] Sifan Wang, Hanwen Wang, and Paris Perdikaris.  Learning the solution operator of parametric partial differential equations with physics-informed deeponets.  _CoRR_ , abs/2103.10974, 2021.  URL https://arxiv.org/abs/2103.10974.
  * Yu et al. [2022] Jeremy Yu, Lu Lu, Xuhui Meng, and George Em Karniadakis.  Gradient-enhanced physics-informed neural networks for forward and inverse pde problems.  _Computer Methods in Applied Mechanics and Engineering_ , 393:114823, 2022.  ISSN 0045-7825.  doi: https://doi.org/10.1016/j.cma.2022.114823.  URL https://www.sciencedirect.com/science/article/pii/S0045782522001438.

##  Appendix A Derivation of primal-dual optimality conditions for dynamical
OT problem

The primal-dual analysis is a standard technique in the optimization
literature such as in analyzing certain semidefinite programs [Chen and Yang,
2021]. Recall the Benamou-Brenier fluid dynamics formulation of the static
optimal transport problem

$\displaystyle\min_{(\mu,\mathbf{v})}\int_{0}^{1}\int_{\mathbb{R}^{d}}{1\over
2}||\mathbf{v}(x,t)||_{2}^{2}\ \mu(x,t)\ \mathrm{d}x\ \mathrm{d}t$ (22)
$\displaystyle\mbox{subject to}\ \
\partial_{t}\mu+\mbox{div}(\mu\mathbf{v})=0,$ (23)
$\displaystyle\mu(\cdot,0)=\mu_{0},\ \ \mu(\cdot,1)=\mu_{1}.$ (24)

Here, equation (23) is referred to as the _continuity equation_ (CE),
preserving the unit mass of the density flow $\mu_{t}=\mu(\cdot,t)$. We write
the Lagrangian function for any flow $(\mu_{t})_{t\in[0,1]}$ initializing from
$\mu_{0}$ and terminating at $\mu_{1}$ as

$\begin{gathered}L(\mu,\mathbf{v},u)=\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[{1\over
2}\|\mathbf{v}\|_{2}^{2}\mu+\left(\partial_{t}\mu+\mbox{div}(\mu\mathbf{v})\right)u\right]\;\mathrm{d}x\;\mathrm{d}t,\end{gathered}$
(25)

where $u:=u(x,t)$ is the dual variable for (CE). To find the optimal solution
$\mu^{*}$ for the minimum kinetic energy (22), we study the saddle point
optimization problem

$\min_{(\mu,\mathbf{v})\in\text{(CE)}}\max_{u}L(\mu,\mathbf{v},u),$ (26)

where the minimization over $(\mu,\mathbf{v})$ runs over all flows satisfying
(CE) such that $\mu(\cdot,0)=\mu_{0}$ and $\mu(\cdot,1)=\mu_{1}$. Note that if
$\mu\notin\text{(CE)}$, then by scaling with arbitrarily large constant, we
see that

$\max_{u}\int_{0}^{1}\int_{\mathbb{R}^{d}}\left(\partial_{t}\mu+\mbox{div}(\mu\mathbf{v})\right)u\;\mathrm{d}x\;\mathrm{d}t=+\infty.$
(27)

Thus,

$\displaystyle\min_{(\mu,\mathbf{v})\in\text{(CE)}}\int_{0}^{1}\int_{\mathbb{R}^{d}}{1\over
2}||\mathbf{v}||_{2}^{2}\mu\ dx\ dt=$
$\displaystyle\min_{(\mu,\mathbf{v})}\max_{u}L(\mu,\mathbf{v},u)$ (28)
$\displaystyle\geqslant$
$\displaystyle\max_{u}\min_{(\mu,\mathbf{v})}L(\mu,\mathbf{v},u),$ (29)

where the minimization over $(\mu,\mathbf{v})$ is unconstrained. Using
integration-by-parts and suitable decay for vanishing boundary as
$\|x\|_{2}\to\infty$, we have

$\displaystyle L(\mu,\mathbf{v},u)=$
$\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[{1\over
2}\|\mathbf{v}\|_{2}^{2}\mu-\mu\partial_{t}u-\langle v,\nabla
u\rangle\mu\right]\;\mathrm{d}x\;\mathrm{d}t$
$\displaystyle\qquad+\int_{\mathbb{R}^{d}}\left[\mu(\cdot,1)u(\cdot,1)-\mu(\cdot,0)u(\cdot,0)\right]\;\mathrm{d}x.$

Now, we fix $\mu$ and $u$, and minimize $L(\mu,\mathbf{v},u)$ over
$\mathbf{v}$. The optimal velocity vector is $\mathbf{v}^{*}=\nabla u$, and we
have

$\max_{u}\min_{\mu}L(\mu,\mathbf{v}^{*},u)=\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[-\left({1\over
2}\|\nabla
u\|_{2}^{2}+\partial_{t}u\right)\mu\right]\;\mathrm{d}x\;\mathrm{d}t+\int_{\mathbb{R}^{d}}\left[u(\cdot,1)\mu_{1}-u(\cdot,0)\mu_{0}\right]\;\mathrm{d}x,$
(30)

for any flow $\mu_{t}$ satisfying the boundary conditions
$\mu(\cdot,0)=\mu_{0}$ and $\mu(\cdot,1)=\mu_{1}$. If ${1\over 2}\|\nabla
u\|_{2}^{2}+\partial_{t}u\neq 0$, then by the same scaling argument above, we
have

$\min_{\mu}\int_{0}^{1}\int_{\mathbb{R}^{d}}\left[-\left({1\over 2}\|\nabla
u\|_{2}^{2}+\partial_{t}u\right)\mu\right]\;\mathrm{d}x\;\mathrm{d}t=-\infty$
(31)

because $\mu$ is unconstrained (except for the boundary conditions). Then we
deduce that

$\min_{(\mu,\mathbf{v})\in\text{(CE)}}\int_{0}^{1}\int_{\mathbb{R}^{d}}{1\over
2}||\mathbf{v}||_{2}^{2}\mu\geqslant\max_{u\in\text{(HJ)}}\left\\{\int_{\mathbb{R}^{d}}u(\cdot,1)\mu_{1}-\int_{\mathbb{R}^{d}}u(\cdot,0)\mu_{0}\right\\},$
(32)

where $u\in\text{(HJ)}$ means that $u$ solves the _Hamilton-Jacobi equation_
(HJ)

$\partial_{t}u+{1\over 2}\|\nabla u\|_{2}^{2}=0.$ (33)

From (32), we see that the duality gap is non-negative, and it is equal to
zero if and only if $(\mu^{*},u^{*})$ solves the following system of PDEs

$\left\\{\begin{gathered}\partial_{t}\mu+\mbox{div}(\mu\nabla u)=0,\ \
\partial_{t}u+{1\over 2}\|\nabla u\|_{2}^{2}=0,\\\ \mu(\cdot,0)=\mu_{0},\ \
\mu(\cdot,1)=\mu_{1}.\end{gathered}\right.$ (34)

PDEs in (34) are referred to as the Karush–Kuhn–Tucker (KKT) conditions for
the Wasserstein geodesic problem.

##  Appendix B Metric geometry structure of the Wasserstein space and geodesic

In this section, we review some basic facts on metric geometry properties of
the Wasserstein space and geodesic. We first discuss the general metric space
$(X,d)$, and then specialize to the Wasserstein (metric) space $({\cal
P}_{p}(\mathbb{R}^{d}),W_{p})$ for $p\geqslant 1$. Furthermore, we connect to
the fluid dynamic formulation of optimal transport. Most of the materials are
based on the reference books [Burage et al., 2001, Ambrosio et al., 2008,
Santambrogio, 2015].

###  B.1 General metric space

######  Definition B.1 (Absolutely continuous curve).

Let $(X,d)$ be a metric space. A curve $\omega:[0,1]\to X$ is _absolutely
continuous_ if there is a function $g\in L^{1}([0,1])$ such that for all
$t_{0}<t_{1}$, we have

$d(\omega(t_{0}),\omega(t_{1}))\leqslant\int_{t_{0}}^{t_{1}}g(\tau)\,\mathrm{d}\tau.$
(35)

Such curves are denoted by $\text{AC}(X)$.

######  Definition B.2 (Metric derivative).

If $\omega:[0,1]\to X$ is a curve in a metric space $(X,d)$, the _metric
derivative_ of $\omega$ at time $t$ is defined as

$|\omega^{\prime}|(t):=\lim_{h\to 0}{d(\omega(t+h),\omega(t))\over|h|},$ (36)

if the limit exists.

The following theorem generalizes the classical Rademacher theorem from a
Euclidean space into any metric space in terms of the metric derivative.

######  Theorem B.3 (Rademacher).

If $\omega:[0,1]\to X$ is Lipschitz continuous, then the metric derivative
$|\omega^{\prime}|(t)$ exists for almost every $t\in[0,1]$. In addition, for
any $0\leqslant t<s\leqslant 1$, we have

$d(\omega(t),\omega(s))\leqslant\int_{t}^{s}|\omega^{\prime}|(\tau)\,\mathrm{d}\tau.$
(37)

Theorem B.3 tells us that absolutely continuous curve $\omega$ has a metric
derivative well-defined almost everywhere, and the “length" of the curve
$\omega$ is bounded by the integral of the metric derivative. Thus, a natural
definition of the length of a curve in a general metric space is to take the
best approximation over all possible meshes.

######  Definition B.4 (Curve length).

For a curve $\omega:[0,1]\to X$, we define its _length_ as

$\text{Length}(\omega):=\sup\left\\{\sum_{k=0}^{n-1}d(\omega(t_{k}),\omega(t_{k+1})):n\geqslant
1,0=t_{0}<t_{1}<\ldots<t_{n}=1\right\\}.$ (38)

Note that if $\omega\in\text{AC}(X)$, then

$d(\omega(t_{k}),\omega(t_{k+1}))\leqslant\int_{t_{k}}^{t_{k+1}}g(\tau)\,\mathrm{d}\tau$
(39)

so that

$\text{Length}(\omega)\leqslant\int_{0}^{1}g(\tau)\,\mathrm{d}\tau<\infty,$
(40)

i.e., the curve $\omega$ is of bounded variation.

######  Lemma B.5.

If $\omega\in\text{AC}(X)$, then

$\text{Length}(\omega)=\int_{0}^{1}|\omega^{\prime}|(\tau)\,\mathrm{d}\tau.$
(41)

######  Definition B.6 (Length space and geodesic space).

Let $\omega:[0,1]\to X$ be a curve in $(X,d)$.

  1. 1.

The space $(X,d)$ is a _length space_ if

$d(x,y)=\inf\left\\{\text{Length}(\omega):\omega(0)=x,\omega(1)=y,\omega\in\text{AC}(X)\right\\}.$
(42)

  2. 2.

The space $(X,d)$ is a _geodesic space_ if

$d(x,y)=\min\left\\{\text{Length}(\omega):\omega(0)=x,\omega(1)=y,\omega\in\text{AC}(X)\right\\}.$
(43)

######  Definition B.7 (Geodesic).

Let $(X,d)$ be a length space.

  1. 1.

A curve $\omega:[0,1]\to X$ is said to be a _constant-speed geodesic_ between
$\omega(0)$ and $\omega(1)$ if

$d(\omega(t),\omega(s))=|t-s|\cdot d(\omega(0),\omega(1)),$ (44)

for any $t,s\in[0,1]$.

  2. 2.

If $(X,d)$ is further a geodesic space, a curve $\omega:[0,1]\to X$ is said to
be a _geodesic_ between $x_{0}\in X$ and $x_{1}\in X$ if it minimizes the
length among all possible curves such that $\omega(0)=x_{0}$ and
$\omega(1)=x_{1}$.

Note that in a geodesic space $(X,d)$, a constant-speed geodesic is indeed a
geodesic. In addition, we have the following equivalent characterization of
the geodesic in a geodesic space.

######  Lemma B.8.

Let $(X,d)$ be a geodesic space, $p>1$, and $\omega:[0,1]\to X$ a curve
connecting $x_{0}$ and $x_{1}$. Then the followings are equivalent.

  1. 1.

$\omega$ is a constant-speed geodesic.

  2. 2.

$\omega\in\text{AC}(X)$ such that for almost every $t\in[0,1]$, we have

$|\omega^{\prime}|(t)=d(\omega(0),\omega(1)).$ (45)

  3. 3.

$\omega$ solves

$\min\left\\{\int_{0}^{1}|\tilde{\omega}^{\prime}|^{p}\,\mathrm{d}t:\tilde{\omega}(0)=x_{0},\tilde{\omega}(1)=x_{1}\right\\}.$
(46)

###  B.2 Wasserstein space

Since the Wasserstein space $({\cal P}_{p}(\mathbb{R}^{d}),W_{p})$ for
$p\geqslant 1$ is a metric space, the following definition specializes
Definition B.2 to the Wasserstein metric derivative.

######  Definition B.9 (Wasserstein metric derivative).

Let $\\{\mu_{t}\\}_{t\in[0,1]}$ be an absolutely continuous curve in the
Wasserstein (metric) space $({\cal P}_{p}(\mathbb{R}^{d}),W_{p})$. Then the
_metric derivative_ at time $t$ of the curve $t\mapsto\mu_{t}$ with respect to
$W_{p}$ is defined as

$|\mu^{\prime}|_{p}(t):=\lim_{h\to 0}{W_{p}(\mu_{t+h},\mu_{t})\over|h|}.$ (47)

For $p=2$, we write $|\mu^{\prime}|(t):=|\mu^{\prime}|_{2}(t)$.

In the rest of this section, we consider probability measures $\mu_{t}$ that
are absolutely continuous with respect to the Lebesgue measure on
$\mathbb{R}^{d}$ and we use $\mu_{t}$ denote the probability measure, as well
as its density, when the context is clear.

######  Theorem B.10.

Let $p>1$ and assume $\Omega\in\mathbb{R}^{d}$ is compact.

Part 1. If $\\{\mu_{t}\\}_{t\in[0,1]}$ is an absolutely continuous curve in
$W_{p}(\Omega)$, then for almost every $t\in[0,1]$, there is a velocity vector
field $\mathbf{v}_{t}\in L^{p}(\mu_{t})$ such that

  1. 1.

$\mu_{t}$ is a weak solution of the continuity equation
$\partial_{t}\mu_{t}+\mbox{div}(\mu_{t}\mathbf{v}_{t})=0$ in the sense of
distributions (cf. the definition in (53) below);

  2. 2.

for almost every $t\in[0,1]$, we have

$\|\mathbf{v}_{t}\|_{L^{p}(\mu_{t})}\leqslant|\mu^{\prime}|_{p}(t),$ (48)

where
$\|\mathbf{v}_{t}\|_{L^{p}(\mu_{t})}^{p}=\int_{\Omega}\|\mathbf{v}_{t}\|_{2}^{p}\,\mathrm{d}\mu_{t}$.

Part 2. Conversely, if $\\{\mu_{t}\\}_{t\in[0,1]}$ are probability measures in
${\cal P}_{p}(\Omega)$, and for each $t\in[0,1]$ we suppose $\mathbf{v}_{t}\in
L^{p}(\mu_{t})$ and
$\int_{0}^{1}\|\mathbf{v}_{t}\|_{L^{p}(\mu)}\,\mathrm{d}t<\infty$ such that
$(\mu_{t},\mathbf{v}_{t})$ solves the continuity equation, then we have

  1. 1.

$\\{\mu_{t}\\}_{t\in[0,1]}$ is an absolutely continuous curve in $({\cal
P}_{p}(\mathbb{R}^{d}),W_{p})$;

  2. 2.

for almost every $t\in[0,1]$,

$|\mu^{\prime}|_{p}(t)\leqslant\|\mathbf{v}_{t}\|_{L^{p}(\mu_{t})}.$ (49)

As an immediate corollary, we have the following dynamical representation of
the Wasserstein metric derivative.

######  Corollary B.11.

If $\\{\mu_{t}\\}_{t\in[0,1]}$ is an absolutely continuous curve in $({\cal
P}_{p}(\mathbb{R}^{d}),W_{p})$, then the velocity vector field
$\mathbf{v}_{t}$ given in Part 1 of Theorem B.10 must satisfy

$\|\mathbf{v}_{t}\|_{L^{p}(\mu_{t})}=|\mu^{\prime}|_{p}(t).$ (50)

Corollary B.11 suggests that $\mathbf{v}_{t}$ can be viewed as the _tangent
vector field_ of the curve $\\{\mu_{t}\\}_{t\in[0,1]}$ at time point $t$.
Moreover, Corollary B.11 suggests the following (Euclidean) gradient flow for
tracking particles in $\mathbb{R}^{d}$: let $y(t):=y_{x}(t)$ be the trajectory
starting from $x\in\mathbb{R}^{d}$ (i.e., $y(0)=x$) that evolves according the
ordinary differential equation (ODE)

${\mathrm{d}\over\mathrm{d}t}y(t)=\mathbf{v}_{t}(y(t)).$ (51)

The dynamical system (51) defines a flow $Y_{t}:\Omega\to\Omega$ of vector
field $\mathbf{v}_{t}$ on $\Omega$ via

$Y_{t}(x)=y(t).$ (52)

Then, it is straightforward to check that the pushforward measure flow
$\mu_{t}:=(Y_{t})_{\sharp}\mu_{0}$ and the chosen velocity vector field
$\mathbf{v}_{t}$ in the ODE (51) is a weak solution of the continuity equation
$\partial_{t}\mu_{t}+\mbox{div}(\mu_{t}\mathbf{v}_{t})=0$ in the sense that

${\mathrm{d}\over\mathrm{d}t}\int_{\Omega}\phi\,\mathrm{d}t=\int_{\Omega}\langle\nabla\phi,\mathbf{v}_{t}\rangle\,\mathrm{d}\mu_{t},$
(53)

for any ${\cal C}^{1}$ function $\phi:\Omega\to\mathbb{R}$ with compact
support.

######  Theorem B.12 (Constant-speed Wasserstein geodesic).

Let $\Omega\in\mathbb{R}^{d}$ be a convex subset and $\mu,\nu\in{\cal
P}_{p}(\Omega)$ for some $p>1$. Let $\gamma$ be an optimal transport plan
under the cost function $\|x-y\|_{p}^{p}$. Define

$\displaystyle\pi_{t}:\Omega\times\Omega\to\Omega,$
$\displaystyle\pi_{t}(x,y)=(1-t)x+ty,$

as the linear interpolation between $x$ and $y$ in $\Omega$. Then, the curve
$\mu_{t}=(\pi_{t})_{\sharp}\gamma$ is a constant-speed geodesic in $({\cal
P}_{p}(\mathbb{R}^{d}),W_{p})$ connecting $\mu_{0}=\mu$ and $\mu_{1}=\nu$.

If $\mu$ has a density with respect to the Lebesgue measure on
$\mathbb{R}^{d}$, then there is an optimal transport map $T$ from $\mu$ to
$\nu$ [Brenier, 1991]. According to Theorem B.12, we obtain _McCann’s
interpolation_ [McCann, 1997] in the Wasserstein space as

$\mu_{t}=[(1-t)\text{id}+tT]_{\sharp}\mu,$ (54)

which is a constant-speed geodesic in $({\cal P}_{p}(\mathbb{R}^{d}),W_{p})$.
id is the identity function in $\mathbb{R}^{d}$.

To sum up, we collect the following facts about the geodesic structure and
dynamical formulation of the OT problem. Let $p>1$, and
$\Omega\subset\mathbb{R}^{d}$ be a convex subset (either compact or have no
mass escaping at infinity).

  1. 1.

The metric space $({\cal P}_{p}(\Omega),W_{p})$ is a geodesic space.

  2. 2.

For $\mu,\nu\in{\cal P}_{p}(\Omega)$, a constant-speed geodesic
$\\{\mu_{t}\\}_{t\in[0,1]}$ in $({\cal P}_{p}(\Omega),W_{p})$ between $\mu$
and $\nu$ (i.e., $\mu_{0}=\mu$ and $\mu_{1}=\nu$) must satisfy
$\mu_{t}\in\text{AC}({\cal P}_{p}(\Omega))$ and

$|\mu^{\prime}|(t)=W_{p}(\mu(0),\mu(1))=W_{p}(\mu,\nu)$ (55)

for almost every $t\in[0,1]$.

  3. 3.

The above $\mu_{t}$ solves

$\min\left\\{\int_{0}^{1}|\tilde{\mu}^{\prime}|^{p}(t)\,\mathrm{d}t:\tilde{\mu}(0)=\mu,\tilde{\mu}(1)=\nu,\tilde{\mu}\in\text{AC}({\cal
P}_{p}(\Omega))\right\\}.$ (56)

  4. 4.

The above $\mu_{t}$ solves the Benamou-Brenier problem

$W_{p}^{p}(\mu,\nu)=\min\left\\{\int_{0}^{1}\|\mathbf{v}_{t}\|_{L^{p}(\tilde{\mu}_{t})}^{p}\,\mathrm{d}t:\tilde{\mu}(0)=\mu,\tilde{\mu}(1)=\nu,\partial_{t}\tilde{\mu}_{t}+\mbox{div}(\tilde{\mu}_{t}\mathbf{v}_{t})=0\right\\},$
(57)

and $\mu_{t}=\mu(\cdot,t)$ defines a constant-speed geodesic in $({\cal
P}_{p}(\Omega),W_{p})$.

##  Appendix C Entropic regularization

Our GeONet is compatible with entropic regularization, which is closely
related to the Schrödinger bridge problem and stochastic control [Chen et al.,
2016]. Specifically, the entropic-regularized GeONet (ER-GeONet) solves the
following fluid dynamic problem:

$\displaystyle\begin{gathered}\min_{(\mu,\mathbf{v})}\int_{0}^{1}\int_{\mathbb{R}^{d}}{1\over
2}||\mathbf{v}(x,t)||_{2}^{2}\ \mu(x,t)\ \mathrm{d}x\ \mathrm{d}t\\\
\qquad\qquad\qquad\qquad\mbox{subject to}\ \
\partial_{t}\mu+\mbox{div}(\mu\mathbf{v})+\varepsilon\Delta\mu=0,\ \
\mu(\cdot,0)=\mu_{0},\ \ \mu(\cdot,1)=\mu_{1}.\end{gathered}$ (60)

Applying the same variational analysis as in the unregularized case
$\varepsilon=0$ (cf. Appendix A), we obtain the KKT conditions for the
optimization (60) as the solution to the following system of PDEs:

$\displaystyle\partial_{t}\mu+\mbox{div}(\mu\nabla u)=$
$\displaystyle-\varepsilon\Delta\mu,$ (61) $\displaystyle\partial_{t}u+{1\over
2}\|\nabla u\|_{2}^{2}=$ $\displaystyle\ \ \varepsilon\Delta u,$ (62)

with the boundary conditions $\mu(\cdot,0)=\mu_{0},\mu(\cdot,1)=\mu_{1}$ for
$\varepsilon>0$. Note that (62) is a parabolic PDE, which has a unique smooth
solution $u^{\varepsilon}$. The term $\varepsilon\Delta u$ effectively
regularizes the (dual) Hamilton-Jacobi equation in (9). In the zero-noise
limit as $\varepsilon\downarrow 0$, the solution of the optimal entropic
interpolating flow (60) converges to solution of the Benamou-Brenier problem
(6) in the sense of the method of vanishing viscosity [Mikami, 2004, Evans,
2010].

##  Appendix D Gradient enhancement

In this section, we fortify the base method we presented in Section 3. We
present gradient enhancement, which is a technique to strengthen any standard
PINN [Yu et al., 2022]. This technique improves efficiency, as fewer data
points are needed to be sampled from $U(\Omega)\times U(0,1)$, and accuracy as
well. We apply gradient enhancement to our proposed neural operator.

The motivation behind gradient enhancement stems minimizing the residual of a
differentiated PDE. We turn our attention to PDEs of the form

$\left\\{\begin{gathered}\mathcal{F}\Big{(}x,t,\partial_{x_{1}}u,\ldots,\partial_{x_{d}}u,\partial_{x_{1}x_{1}}u,\ldots,\partial_{x_{d}x_{d}}u,\ldots,\partial_{t}u,\lambda\Big{)}=0\
\ \ \ \text{on}\ \ \ \ \Omega\times[0,1],\\\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \
\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ u(\cdot,0)=u_{0},\ \ \ u(\cdot,1)=u_{1}\
\ \ \ \ \ \ \ \ \ \ \ \ \text{on}\ \ \ \ \Omega,\ \ \end{gathered}\right.$
(63)

for domain $\Omega\subseteq\mathbb{R}^{d}$, parameter vector $\lambda$, and
boundary conditions $u_{0},u_{1}$. One may differentiate the PDE function
$\mathcal{F}$ with respect to any spatial component to achieve

$\frac{\partial}{\partial
x_{\ell}}\mathcal{F}\Big{(}x,t,\partial_{x_{1}}u,\ldots,\partial_{x_{d}}u,\partial_{x_{1}x_{1}}u,\ldots,\partial_{x_{d}x_{d}}u,\ldots,\partial_{t}u,\lambda\Big{)}=0.$
(64)

The differentiated PDE is additionally equal to $0$, similar to what we see in
a PINN setup. If we substitute a neural network into the differentiated PDE of
(64), what remains is a new residual, just as we saw with the neural network
substituted into the original PDE. Minimizing this new residual in the related
loss function characterizes the gradient enhancement method.

We utilize the same loss function in (18), but we add the additional terms

$\displaystyle\mathcal{L}_{\text{GE,cty},i}\ \ \ $ $\displaystyle=\ \ \
\frac{1}{N}\sum_{\ell=1}^{d}\gamma_{\ell}\Big{|}\Big{|}\ \
\frac{\partial}{\partial x_{\ell}}\Big{(}\frac{\partial}{\partial
t}\mathcal{C}_{\phi,i}+\mbox{div}(\mathcal{C}_{\phi,i}\nabla\mathcal{H}_{\psi,i})\Big{)}\
\ \Big{|}\Big{|}_{L^{2}(\Omega\times(0,1))}^{2},$ (65)
$\displaystyle\mathcal{L}_{\text{GE, HJ},i}\ \ \ $ $\displaystyle=\ \ \
\frac{1}{N}\sum_{\ell=1}^{d}\omega_{\ell}\Big{|}\Big{|}\ \
\frac{\partial}{\partial x_{\ell}}\Big{(}\frac{\partial}{\partial
t}\mathcal{H}_{\psi,i}+\frac{1}{2}||\nabla\mathcal{H}_{\psi,i}||_{2}^{2}\Big{)}\
\ \Big{|}\Big{|}_{L^{2}(\Omega\times(0,1))}^{2},$ (66)

where the variables and neural networks that also appeared in (18) are the
same. Here $\gamma_{\ell}$ and $\omega_{\ell}$ are positive weights. The
summation is taken in order to account for gradient enhancement of each
spatial component of $x\in\Omega$.

##  Appendix E DeepONets

A challenge resides in solving the previous risk minimization problem over
numerous instances of data. This challenge may be conciliated by instituting
the novel architecture of the DeepONet that learns a general nonlinear
operator, where one (or a pair of) neural network(s) encode(s) the input and
another encodes the collocation samples. This architecture originates as a
fine equivalence to the universal approximation theorem for operators.

General DeepONet. A general operator $G^{\dagger}$ may be approximated by an
unstacked DeepONet [Chen and Chen, 1995, Lu et al., 2021]

$G^{\dagger}(u_{0})(x,t)\approx\sum_{k=1}^{p}\mathcal{B}_{k}\big{(}u_{0}(x_{1}),\ldots,u_{0}(x_{m}),\theta\big{)}\cdot\mathcal{T}_{k}(x,t,\xi),$
(67)

where $\mathcal{B}_{k},\mathcal{T}_{k}$ are scalar elements of output of
neural networks $\mathcal{B},\mathcal{T}$, and $p$ is a constant denoting the
number of such elements. We take $\mathcal{B}$ and $\mathcal{T}$ to be
artificial neural networks parameterized by $\theta,\xi$ respectively.
$\mathcal{B},\mathcal{T}$ are known as the branch and trunk networks
respectively. $u_{0}$ is the initial function in which the operator is
applied, evaluated at distinct locations $x_{1},\ldots,x_{m}$ for branch
input. $(x,t)$ is any arbitrary point in space and time in which
$G^{\dagger}(u_{0})$ may be evaluated.

Enhanced DeepONet. The above framework is restricted to one initial input
function $u_{0}$. We turn our attention to the enhanced DeepONet, a DeepONet
styled to act upon dual initial conditions Tan and Chen [2022]. Our true
operator $\Gamma^{\dagger}$ may be approximated using a second neural network
encoder for input $u_{1}$,

$\Gamma^{\dagger}(u_{0},u_{1})(x,t)\approx\sum_{k=1}^{p}\mathcal{B}^{0}_{k}\big{(}u_{0}(x_{1}),\ldots,u_{0}(x_{m}),\theta^{0}\big{)}\cdot\mathcal{B}^{1}_{k}\big{(}u_{1}(x_{1}),\ldots,u_{1}(x_{m}),\theta^{1}\big{)}\cdot\mathcal{T}_{k}(x,t,\xi).$
(68)

Physics-informed DeepONet. The enhanced DeepONet may be substituted into any
physics-informed framework, such as that of equation (13), taking place of the
PDE solution value in the empirical loss to be minimized Wang et al. [2021].
Generalization of the trained DeepONet permits any solution to the PDEs to be
evaluated instantaneously given the appropriate input function(s).

##  Appendix F DeepONet differentiation

We found it effective to differentiate our trunk networks individually, and
subsequently form the DeepONets afterwards. It can be noted the branch
networks take no $(x,t)$ input, and so no branch derivatives are taken. The
product rule of calculus allows the derivatives in the continuity physics-
informed term to be reformed, which can be stated as

$\begin{gathered}\mbox{div}(\mathcal{C}_{\phi}^{j}\nabla\mathcal{H}_{\psi}^{j})=\sum_{j=1}^{d}\Big{\\{}(\sum_{k=1}^{p}\mathcal{B}_{k}^{0,\text{cty}}\mathcal{B}_{k}^{1,\text{cty}}\partial_{x_{j}}\mathcal{T}_{k}^{\text{cty}})(\sum_{k=1}^{p}\mathcal{B}_{k}^{0,\text{HJ}}\mathcal{B}_{k}^{1,\text{HJ}}\partial_{x_{j}}\mathcal{T}_{k}^{\text{HJ}})\\\
\qquad\qquad\qquad\qquad\qquad\qquad+(\sum_{k=1}^{p}\mathcal{B}_{k}^{0,\text{cty}}\mathcal{B}_{k}^{1,\text{cty}}\mathcal{T}_{k}^{\text{cty}})(\sum_{k=1}^{p}\mathcal{B}_{k}^{0,\text{HJ}}\mathcal{B}_{k}^{1,\text{HJ}}\partial_{x_{j}}^{2}\mathcal{T}_{k}^{\text{HJ}})\Big{\\}},\end{gathered}$
(69)

where the above trunk network derivatives are evaluated at
$(x^{j},t^{j}),x^{j}=(x_{1}^{j},\ldots,x_{d}^{j})\in\Omega\subseteq\mathbb{R}^{d}$.
Similarly, the norm-gradient term in the Hamilton-Jacobi physics-informed loss
can be reformulated

$\displaystyle||\nabla\mathcal{H}_{\psi}^{j}||_{2}^{2}=\sum_{j=1}^{d}(\sum_{k=1}^{p}\mathcal{B}_{k}^{0,\text{HJ}}\mathcal{B}_{k}^{1,\text{HJ}}\partial_{x_{j}}\mathcal{T}_{k}^{\text{HJ}})^{2}.$
(70)

We reiterate the above derivatives are computed with automatic
differentiation.

##  Appendix G Additional simulation result for Gaussian mixtures

Here we present additional simulation result for Gaussian mixtures with
$k_{0}=k_{1}=5$ and $\pi_{i}=0.2$ for all $i$, with the same loss
coefficients. We choose $u_{i}\in[1.1,3.9]^{2}$,
$\sigma_{0,i}^{2},\sigma_{1,i}^{2}\in[0.4,0.8]$, and covariance
$\sigma_{01,i}\in[-0.4,0.4]$. The MSE of GeONet on testing pairs is shown in
Table 4.

Table 4: MSE of GeONet on 50 test data of Gaussian mixtures over a $50\times 50$ mesh. All values are multiplied by $10^{-3}$. We report the means and standard deviations of the MSE. Training is done with $k_{0}=k_{1}=5$, $\pi_{i}=0.2$ for all $i$, with the same loss coefficients. Training data at the boundaries has resolution $30\times 30$.  | GeONet error for Gaussian mixtures
---|---
Number of Gaussians | $\bm{t=0}$ | $\bm{t=0.25}$ | $\bm{t=0.5}$ | $\bm{t=0.75}$ | $\bm{t=1}$
Identity $k_{0}=k_{1}=5$ | $0.23\pm 0.23$ | $1.70\pm 0.67$ | $2.60\pm 1.10$ | $1.80\pm 0.74$ | $0.22\pm 0.21$
Generic $k_{0}=k_{1}=5$ | $0.22\pm 0.15$ | $1.70\pm 0.92$ | $2.70\pm 1.50$ | $1.70\pm 0.76$ | $0.20\pm 0.13$

##  Appendix H Hyperparameter settings and training details

We discuss training characteristics of GeONet. We provide details of the base
method. An unmodified Adam optimizer was chosen for all branch, trunk neural
networks. The below widths are held for each hidden layer. The activation
listed is for each hidden layer in all six neural networks. We found constant
learning rate to be successful, but one can consider modifying this.
Coefficients $\alpha_{1},\alpha_{2},\beta_{0},\beta_{1}$ were chosen after
experimentation on sub-datasets. Training is done on a NVIDIA Tesla P100 GPU.

Table 5: Architecture and training details in our experimental Section 4. Gaussian mixture details pertain to Table 4. Hyperparameter/Characteristic | Gaussian mixtures | CIFAR-10 data
---|---|---
No. of initial conditions $(\mu_{0},\mu_{1})$ | 1,500 | 600
$L$ (no. of collocations per epoch) | 900 | 1,024
Branch width | 180 | 300
Branch depth | 5 | 4
Trunk width | 120 | 120
Trunk depth | 7 | 6
$p$ (dimension of branch, trunk outputs) | 120 | 250
Number of parameters | 1,461,120 | 2,821,260
Batch size | 9,000 | 10,240
Number of batches | 150 | 60
Activation | tanh | gelu
Learning rate | $5\times 10^{-5}$ | $4\times 10^{-5}$
Time per epoch |  $\sim 22$ sec |  $\sim 12$ sec
Total number of epochs | $\sim 2,500$ | $\sim 4,500$
Final training time |  $\sim 15$ hrs |  $\sim 15$ hrs
Final training loss | $\sim 4\times 10^{-6}$ | $\sim 1\times 10^{-4}$
$\alpha_{1},\alpha_{2},\beta_{0},\beta_{1}$ | $30,30,1,1$ | $10,10,1,1$