#  Invertible normalizing flow neural networks by JKO scheme

Chen Xu  H. Milton Stewart School of Industrial and Systems Engineering,
Georgia Institute of Technology. Xiuyuan Cheng  Department of Mathematics,
Duke University Yao Xie  H. Milton Stewart School of Industrial and Systems
Engineering, Georgia Institute of Technology.

###### Abstract

Normalizing flow is a class of deep generative models for efficient sampling
and density estimation. In practice, the flow often appears as a chain of
invertible neural network blocks; to facilitate training, existing works have
regularized flow trajectories and designed special network architectures. The
current paper develops a neural ODE flow network inspired by the Jordan-
Kinderleherer-Otto (JKO) scheme, which allows efficient block-wise training of
the residual blocks and avoids inner loops of score matching or variational
learning. As the JKO scheme unfolds the dynamic of gradient flow, the proposed
model naturally stacks residual network blocks one-by-one, reducing the memory
load and difficulty of performing end-to-end training of deep flow networks.
We also develop adaptive time reparameterization of the flow network with a
progressive refinement of the trajectory in probability space, which improves
the model training efficiency and accuracy in practice. Using numerical
experiments with synthetic and real data, we show that the proposed JKO-iFlow
model achieves similar or better performance in generating new samples
compared with existing flow and diffusion models at a significantly reduced
computational and memory cost.

##  1 Introduction

a JKO-iFlow

b usual CNF

Figure 1:  Comparison of JKO-iFlow (proposed) and standard Continuous
Normalizing Flow (CNF) flow models, a popular framework in the literature. In
contrast to usual CNF models, JKO-iFlow learns the unique smooth flow
corresponding to the diffusion process and can be trained block-wisely.

Generative models have wide applications in statistics and machine learning to
infer data-generating distributions and to sample from the model distributions
learned from data. In addition to widely used deep generative models such as
variational auto-encoders (VAE) Kingma and Welling (2014, 2019) and generative
adversarial networks (GAN) Goodfellow et al. (2014); Gulrajani et al. (2017),
normalizing flow models Kobyzev et al. (2020) have been popular and with a
great potential. The flow model learns the data distribution via an invertible
mapping $F$ between the data density $p_{X}$ in $\mathbb{R}^{d}$ and the
target standard multivariate Gaussian density $p_{Z}$,
$Z\sim\mathcal{N}(0,I_{d})$. While flow models, once trained, can be utilized
for efficient data sampling and explicit likelihood evaluation, training of
such models is often difficult in practice. To alleviate such difficulties,
many prior works (Dinh et al., 2014, 2017; Kingma and Dhariwal, 2018;
Grathwohl et al., 2019; Behrmann et al., 2019; Onken et al., 2021), among
others, have explored designs of training objectives, network architectures,
and computational techniques.

Among various flow models, continuous normalizing flow (CNF) transports the
data density to a target distribution through continuous dynamics, primarily
neural ODE (Chen et al., 2018) models. CNF models have shown promising
performance on generative tasks Grathwohl et al. (2019); Kobyzev et al.
(2020). However, a known computational challenge of CNF models is model
regularization, primarily due to the non-uniqueness of the flow transport.
Without additional regularization, the trained CNF model may follow a less
regular trajectory in the probability space, cf. Figure 1(b), which may worsen
the generative performance. While regularization of flow models has been
developed using different techniques, including using spectral normalization
Behrmann et al. (2019) and optimal transport Liutkus et al. (2019); Onken et
al. (2021); Finlay et al. (2020); Xu et al. (2022), only using regularization
may not resolve the non-uniqueness of the flow. Besides regularization,
several practical difficulties remain when training CNF models, particularly
the high computational cost: In many settings, CNF models consist of stacked
blocks and each can be complex. End-to-end training of such deep models often
places a high demand on computational resources and memory consumption.

In this work, we propose JKO-iFlow, a normalizing flow network that unfolds
the Wasserstein gradient flow via a neural ODE invertible network inspired by
the JKO-scheme Jordan et al. (1998). The JKO scheme, cf. (5), can be viewed as
a proximal step to unfold the Wasserstein gradient flow to minimize the KL
divergence (relative entropy) between the current density and the equilibrium
density, and it recovers the solution to the Fokker-Planck equation in the
limit of small step size. The proposed JKO-iFlow model thus can be viewed as
trained to learn the unique transport map following the Fokker-Planck equation
which flows from the data distribution toward the normal equilibrium, which
gives a smooth trajectory of density evolution, cf. Figure 1(a). Unlike most
CNF models where all the residual blocks are initialized together and trained
end-to-end, in the JKO-iFlow network, each block in the flow model implements
one step in the JKO scheme and is trained by minimizing an objective of that
block given the trained previous blocks. The block-wise training significantly
reduces memory and computational load. The theoretical need of small step
sizes for the convergence of the JKO scheme does not incur a restriction in
practice when training the flow model, whereby one can use larger step sizes
when adopting numerical integration techniques. We further introduce time
reparameterization with progressive refinement in computing the flow network,
where each block corresponds to a representative point along the density
evolution trajectory in the space of probability measures. In the algorithm,
one can thus determine the number of blocks and adjust the step sizes
adaptively.

The proposed approach is closely related to the diffusion models (Song and
Ermon, 2019; Ho et al., 2020; Block et al., 2020; Song et al., 2021) yet
differs in that it trains a neural ODE model directly without SDE sampling nor
learning of score matching. It also differs from previous works on progressive
training of ResNet generative models Johnson and Zhang (2019); Fan et al.
(2021) in that the model trains an invertible flow mapping and avoids inner
loops of variational learning. We refer to Section 1.1 for more comments on
related works. Empirically, JKO-iFlow yields competitive performance as other
CNF models with significantly less computation. The model is also compatible
with general equilibrium density $p_{Z}$ with parametrized potential $V$,
exemplified by the application to the conditional generation setting where
$p_{Z}$ are Gaussian mixtures Xu et al. (2022).

In summary, the contributions of the work include

$\bullet$ Propose a neural ODE model where each residual block computes a JKO
step, and the training objective can be computed by integrating the ODE on
data samples. The network has a general form, and invertibility can be
satisfied due to the regularity of the optimal pushforward map that minimizes
the objective in each JKO step.

$\bullet$ Develop a block-wise procedure to train the invertible JKO-iFlow
network, which determines the number of blocks adaptively. We also propose
techniques to reparameterize and refine the computed trajectory in probability
space, which improves the efficiency and accuracy of the approach.

$\bullet$ Show that JKO-iFlow greatly reduces memory consumption and
computational cost when achieving competitive or better performance compared
to existing flow and diffusion models.

###  1.1 Related works

For deep generative models, popular approaches include generative adversarial
networks (GAN) (Goodfellow et al., 2014; Gulrajani et al., 2017; Isola et al.,
2017) and variational auto-encoder (VAE)(Kingma and Welling, 2014, 2019).
Apart from known training difficulties (e.g., mode collapse (Salimans et al.,
2016) and posterior collapse (Lucas et al., 2019)), these models do not
provide likelihood or inference of data density. The normalizing flow
framework (Kobyzev et al., 2020) has been extensively developed, including
continuous flow (Grathwohl et al., 2019), Monge-Ampere flow (Zhang et al.,
2018), discrete flow (Chen et al., 2019), graph flow (Liu et al., 2019), etc.
Efforts have been made to develop novel invertible mapping structures (Dinh et
al., 2017; Papamakarios et al., 2017), regularize the flow trajectories
(Finlay et al., 2020; Onken et al., 2021), and extend the use to non-Euclidean
data (Mathieu and Nickel, 2020; Xu et al., 2022). Despite such efforts, the
model and computational challenges of normalizing flow models include
regularization and the large model size when using a large number of residual
blocks, which cannot be determined a priori, and the associated memory and
computational load.

In parallel to continuous normalizing flows which are neural ODE models,
neural SDE models become an emerging tool for generative tasks. Diffusion
process and Langevin dynamics in deep generative models have been studied in
score-based generative models (Song and Ermon, 2019; Ho et al., 2020; Block et
al., 2020; Song et al., 2021) under a different setting. Specifically, these
models estimate the score function (i.e., gradient of the log probability
density with respect to data) of data distribution via neural network
parametrization, which may encounter challenges in learning and sampling of
high dimensional data and call for special techniques (Song and Ermon, 2019).
The recent work of Song et al. (2021) developed reverse-time SDE sampling for
score-based generative models, and adopted the connection to neural ODE to
compute the likelihood; using the same idea of backward SDE, Zhang and Chen
(2021) proposed joint training of forward and backward neural SDEs.
Theoretically, latent diffusion Tzen and Raginsky (2019b, a) was used to
analyze neural SDE models. The current work focuses on a neural ODE model
where the deterministic vector field ${\bf f}(x,t)$ is to be learned from data
following a JKO scheme of the Fokker-Planck equation. Rather than neural SDE,
our approach involves no sampling of SDE trajectories nor learning of the
score function, and it learns an invertible residual network directly. In
contrast, diffusion-based models derive the ODE model from the learned
diffusion model to achieve explicit likelihood computation. For example, Song
et al. (2021) derived neural ODE model from the learned score function of the
diffused data marginal distributions for all $t$. We experimentally obtain
competitive or improved performance against diffusion model on simulated 2D
data and high-dimensional tabular data.

JKO-inspired deep models have been studied in several recent works. (Bunne et
al., 2022) reformulated the JKO step for minimizing an energy function over
convex functions. JKO scheme has also been used to discretize Wasserstein
gradient flow to learn a deep generative model in (Alvarez-Melis et al., 2021;
Mokrov et al., 2021), which adopted input convex neural networks (ICNN) (Amos
et al., 2017). ICNN as a special type of network architecture may have limited
expressiveness (Rout et al., 2022; Korotin et al., 2021). In addition to using
the gradient of ICNN, (Fan et al., 2021) proposed to parametrize the transport
in a JKO step by a residual network but identified difficulty in calculating
the push-forward distribution. The approach in (Fan et al., 2021) also relies
on a variational formulation which requires training an additional network
similar to the discriminator in GAN using inner loops. The idea of progressive
additive learning in training generative ResNet, namely training ResNet block-
wisely by a variational loss, dates back to Johnson and Zhang (2019) under the
GAN framework. Our method trains an invertible neural-ODE flow network that
flows from data density to the normal one and backward, which enables the
computation of model likelihood as in other neural-ODE approaches. The
objective in each JKO step to minimize KL divergence can also be computed
directly without any inner-loop training, cf. Section 4.

For the expressiveness of generating deep models, universal approximation
properties of deep neural networks for representing probability distributions
have been developed in several works. Lee et al. (2017) established
approximation by composition of Barron functions (Barron, 1993); Bailey and
Telgarsky (2018) developed space-filling approach, which was generalized in
Perekrestenko et al. (2020, 2021); Lu and Lu (2020) constructed a deep ReLU
network with guaranteed approximation under integral probability metrics,
using techniques of empirical measures and optimal transport. These results
show that deep neural networks can provably transport one source distribution
to a target one with sufficient model capacity under certain regularity
conditions of the pair of densities. In our proposed flow model, each residual
block is trained to approximate the vector field ${\bf f}(x,t)$ that induces
the Fokker-Planck equation, cf. Section 3.2. Our model potentially leads to a
constructive approximation analysis of neural ODE flow model to generate data
density $p_{X}$.

##  2 Preliminaries

Normalizing flow. A normalizing flow can be mathematically expressed via a
density evolution equation of $\rho(x,t)$ such that $\rho(x,0)=p_{X}$ and as
$t$ increases $\rho(x,t)$ approaches $p_{Z}\sim\mathcal{N}(0,I_{d})$ Tabak and
Vanden-Eijnden (2010). Given an initial distribution $\rho(x,0)$, such a flow
typically is not unique. We consider when the flow is induced by an ODE of
$x(t)$ in $\mathbb{R}^{d}$

$\dot{x}(t)={\bf f}(x(t),t),$ (1)

where $x(0)\sim p_{X}$. The marginal density of $x(t)$ is denoted as $p(x,t)$,
and it evolves according to the continuity equation (Liouville equation) of
(1) written as

$\partial_{t}p+\nabla\cdot(p{\bf f})=0,\quad p(x,0)=p_{X}(x).$ (2)

Ornstein–Uhlenbeck (OU) process. Consider a Langevin dynamic denoted by the
SDE $dX_{t}=-\nabla V(X_{t})dt+\sqrt{2}dW_{t}$, where $V$ is the potential of
the equilibrium density. We focus on the case of normal equilibrium, that is,
$V(x)=|x|^{2}/2$ and then $p_{Z}\propto e^{-V}$. In this case, the process is
known as the (multivariate) OU process. Suppose $X_{0}\sim p_{X}$, and let the
density of $X_{t}$ be $\rho(x,t)$ also denoted as $\rho_{t}(\cdot)$. The
Fokker-Planck equation describes the evolution of $\rho_{t}$ towards the
equilibrium $p_{Z}$ as

$\partial_{t}\rho=\nabla\cdot(\rho\nabla V+\nabla\rho),\quad
V(x):={|x|^{2}}/{2},\quad\rho(x,0)=p_{X}(x).$ (3)

Under generic conditions, $\rho_{t}$ converges to $p_{Z}$ exponentially fast.
For Wasserstein-2 distance and the standard normal $p_{Z}$, classical argument
gives that (take $C=1$ in Eqn (6) of Bolley et al. (2012))

$W_{2}(\rho_{t},p_{Z})\leq e^{-t}W_{2}(\rho_{0},p_{Z}),\quad t>0.$ (4)

JKO scheme. The seminal work Jordan et al. (1998) established a time
discretization scheme of the solution to (3) by the gradient flow to minimize
${\rm KL}(\rho||p_{Z})$ under the Wasserstein-2 metric in probability space.
Denote by $\mathcal{P}$ the space of all probability densities on
$\mathbb{R}^{d}$ with a finite second moment. The JKO scheme computes a
sequence of distributions $p_{k}$, $k=0,1,\cdots,$ starting from
$p_{0}=\rho_{0}\in\mathcal{P}$. With step size $h>0$, the scheme at the $k$-th
step is written as

$p_{k+1}=\arg\min_{\rho\in\mathcal{P}}F[\rho]+\frac{1}{2h}W_{2}^{2}(p_{k},\rho),\quad
F[\rho]:={\rm KL}(\rho||p_{Z}).$ (5)

It was proved in Jordan et al. (1998) that as $h\to 0$, the solution $p_{k}$
converges to the solution $\rho(\cdot,kh)$ of (3) for all $k$, and the
convergence $\rho_{h}(\cdot,t)\to\rho(\cdot,t)$ is strongly in
$L^{1}(\mathbb{R}^{d},(0,T))$ for finite $T$, where $\rho_{h}$ is piece-wise
constant interpolated on $(0,T)$ from $p_{k}$.

##  3 JKO scheme by neural ODE

Given i.i.d. observed data samples $X_{i}\in\mathbb{R}^{d}$, $i=1,\ldots,N$,
drawn from some unknown density $p_{X}$, the goal is to train an invertible
neural network to transports the density $p_{X}$ to an a priori specified
density $p_{Z}$ in $\mathbb{R}^{d}$, where each data sample $X_{i}$ is mapped
to a code $Z_{i}$. A prototypical choice of $p_{Z}$ is the standard
multivariate Gaussian $\mathcal{N}(0,I_{d})$. In this work, we leave the
potential of $p_{Z}$ abstract and denote by $V$, that is, $p_{Z}\propto
e^{-V}$ and $V(x)=|x|^{2}/2$ for normal $p_{Z}$. By a slight abuse of
notation, we denote by $p_{X}$ and $p_{Z}$ both the distributions and the
density functions of data $X$ and code $Z$ respectively.

###  3.1 The objective of JKO step

We are to specify ${\bf f}(x,t)$ in the ODE (1), to be parametrized and
learned by a neural ODE, such that the induced density evolution of $p(x,t)$
converges to $p_{Z}$ as $t$ increases. We start by dividing the time horizon
$[0,T]$ into finite subintervals with step size $h$, let $t_{k}=kh$ and
$I_{k+1}:=[t_{k},t_{k+1})$. Define $p_{k}(x):=p(x,kh)$, namely the density of
$x(t)$ at $t=kh$. The solution of (1) determined by the vector-field ${\bf
f}(x,t)$ on $t\in I_{k+1}$ (assuming the ODE is well-posed (Sideris, 2013))
gives a one-to-one mapping $T_{k+1}$ on $\mathbb{R}^{d}$, s.t.
$T_{k+1}(x(t_{k}))=x(t_{k+1})$ and $T_{k+1}$ transports $p_{k}$ into
$p_{k+1}$, i.e., $(T_{k})_{\\#}p_{k-1}=p_{k}$, where we denote by $T_{\\#}p$
the push-forward of distribution $p$ by $T$, such that
$(T_{\\#}p)(\cdot)=p(T^{-1}(\cdot))$.

Suppose we can find ${\bf f}(\cdot,t)$ on $I_{k+1}$ such that the
corresponding $T_{k+1}$ solves the JKO scheme (5), then with small $h$,
$p_{k}$ approximates the solution to the Fokker-Planck equation 3, which then
flows towards $p_{Z}$. By the Monge formulation of the Wasserstein-2 distance
between $p$ and $q$ as $W_{2}^{2}(p,q)=\min_{T:T_{\\#}p=q}\mathbb{E}_{x\sim
p}\|x-T(x)\|^{2}$, solving for the transported density $p_{k}$ by (5) is
equivalent to solving for the transport $T_{k+1}$ by

$T_{k+1}=\arg\min_{T:\mathbb{R}^{d}\to\mathbb{R}^{d}}F[T]+\frac{1}{2h}\mathbb{E}_{x\sim
p_{k}}\|x-T(x)\|^{2},\quad F[T]={\rm KL}(T_{\\#}p_{k}||p_{Z}).$ (6)

The equivalence between (5) and (6) is proved in Lemma A.1.

Furthermore, the following proposition gives that the value of $F[T]$ can be
computed from ${\bf f}(x,t)$ on $t\in I_{k+1}$ only once $p_{k}$ is determined
by ${\bf f}(x,t)$ for $t\leq t_{k}$. The counterpart for convex function-based
parametrization of $T_{k}$ was given in Theorem 1 of (Mokrov et al., 2021),
where the computation using the change-of-variable differs as we adopt an
invertible neural ODE approach here. The proof is left to Appendix A.

######  Proposition 3.1.

Given $p_{k}$, up to a constant $c$ independent from ${\bf f}(x,t)$ on $t\in
I_{k+1}$,

${\rm KL}(T_{\\#}p_{k}||p_{Z})=\mathbb{E}_{x(t_{k})\sim
p_{k}}\left(V(x(t_{k+1}))-\int_{t_{k}}^{t_{k+1}}\nabla\cdot\mathbf{f}(x(s),s)ds\right)+c.$
(7)

By Proposition 3.1, the minimization (6) is equivalent to

$\min_{\\{{\bf f}(x,t)\\}_{t\in I_{k+1}}}\mathbb{E}_{x(t_{k})\sim
p_{k}}\left(V(x(t_{k+1}))-\int_{t_{k}}^{t_{k+1}}\nabla\cdot\mathbf{f}(x(s),s)ds+\frac{1}{2h}\|x(t_{k+1})-x(t_{k})\|^{2}\right),$
(8)

where $x(t_{k+1})=x(t_{k})+\int_{t_{k}}^{t_{k+1}}\mathbf{f}(x(s),s)ds$. Taking
a neural ODE approach, we parametrize $\\{{\bf f}(x,t)\\}_{t\in I_{k+1}}$ as a
residual block with parameter $\theta_{k+1}$, and then (8) is reduced to
minimizing over $\theta_{k+1}$. This leads to block-wise learning algorithm to
be introduced in Section 4, where we furtherly allow the step-size $h$ to vary
for different $k$ as well.

###  3.2 Infinitesimal optimal ${\bf f}(x,t)$

In each JKO step of (8), let $p=p_{k}$ denote the current density, $q=p_{Z}$
be the target equilibrium density. In this subsection, we show that the
optimal ${\bf f}$ in (8) with small $h$ reveals the difference between score
functions between target and current densities. Thus minimizing the objective
(8) searches for a neural network parametrization of the score function
$\nabla\log\rho_{t}$ implicitly, in contrast to diffusion-based models which
learn the score function explicitly (Ho et al., 2020; Song et al., 2021), e.g.
via denoising score matching.

Consider general equilibrium distribution $q$ with a differentiable potential
$V$. To analyze the optimal pushforward mapping in the small $h$ limit, we
shift the time interval $[kh,(k+1)h]$ to be $[0,h]$ to simplify notation. Then
(8) is reduced to

$\min_{\\{{\bf f}(x,t)\\}_{t\in[0,h)}}\mathbb{E}_{x(0)\sim
p}\left(V(x(h))-\int_{0}^{h}\nabla\cdot\mathbf{f}(x(s),s)ds+\frac{1}{2h}\|x(h)-x(0)\|^{2}\right),$
(9)

where $x(h)=x(0)+\int_{0}^{h}\mathbf{f}(x(s),s)ds$. In the limit of $h\to 0+$,
formally, $x(h)-x(0)=h{\bf f}(x(0),0)+O(h^{2})$, and suppose $V$ of $q$ is
$C^{2}$, $V(x(h))=V(x(0))+h\nabla V(x(0))\cdot{\bf f}(x(0),0)+O(h^{2})$. For
any differentiable density $\rho$, the (Stein) score function is defined as
${\bf s}_{\rho}=\nabla\log\rho$, and we have $\nabla V=-{\bf s}_{q}$. Taking
the formal expansion of orders of $h$, the objective in (9) is written as

$\mathbb{E}_{x\sim p}\left(V(x)+h\left(-{\bf s}_{q}(x)\cdot{\bf
f}(x,0)-\nabla\cdot\mathbf{f}(x,0)+\frac{1}{2}\|{\bf
f}(x,0)\|^{2}\right)+O(h^{2})\right).$ (10)

Note that $\mathbb{E}_{x\sim p}V(x)$ is independent of ${\bf f}(x,t)$, and the
$O(h)$ order term in (10) is over ${\bf f}(x,0)$ only, thus the minimization
of the leading term is equivalent to

$\min_{{\bf f}(\cdot)={\bf f}(\cdot,0)}\mathbb{E}_{x\sim p}\left(-T_{q}{\bf
f}+\frac{1}{2}\|{\bf f}\|^{2}\right),\quad T_{q}{\bf f}:={\bf s}_{q}\cdot{\bf
f}+\nabla\cdot{\bf f},$ (11)

where $T_{q}$ is known as the Stein operator (Stein, 1972). The $T_{q}{\bf f}$
in (11) echoes that the derivative of KL divergence with respect to transport
map gives Stein operator (Liu and Wang, 2016). The Wasserstein-2
regularization gives an $L^{2}$ regularization in (11). Let $L^{2}(p)$ be the
$L^{2}$ space on $(\mathbb{R}^{d},p(x)dx)$, and for vector field ${\bf v}$ on
$\mathbb{R}^{d}$, ${\bf v}\in L^{2}(p)$ if $\int|{\bf
v}(x)|^{2}p(x)dx<\infty$. One can verify that, when both ${\bf s}_{p}$ and
${\bf s}_{q}$ are in $L^{2}(p)$, the minimizer of (11) is

${\bf f}^{*}(\cdot,0)={\bf s}_{q}-{\bf s}_{p}.$

This shows that the infinitesimal optimal ${\bf f}(x,t)$ equals the difference
of the score functions of the equilibrium and the current density.

###  3.3 Invertibility of flow model and expressiveness

At time $t$ the current density of $x(t)$ is $\rho_{t}$, the analysis in
Section 3.2 implies that the optimal vector field ${\bf f}(x,t)$ has the
expression as

${\bf f}(x,t)={\bf s}_{q}-{\bf s}_{\rho_{t}}=-\nabla V-\nabla\log\rho_{t}.$
(12)

With this ${\bf f}(x,t)$, the Liouville equation (2) coincides with the
Fokker-Planck equation (3). This is consistent with that JKO scheme with small
$h$ recovers the solution to the Fokker-Planck equation. Under proper
regularity condition of $V$ and the initial density $\rho_{0}$, the r.h.s. of
(12) is also regular over space and time. This leads to two consequences, in
approximation and in learning: Approximation-wise, the regularity of ${\bf
f}(x,t)$ allows to construct a $k$-th residual block in the flow network to
approximate $\\{{\bf f}(x,t)\\}_{t\in I_{k}}$ when there is sufficient model
capacity, by classical universal approximation theory of shallow networks
(Barron, 1993; Yarotsky, 2017). The JKO-iFlow model proposed in this work
suggests a constructive proof of the expressiveness of the invertible neural
ODE model to generate any sufficiently regular density $p_{X}$, which we
further discuss in the last section.

For learning, when properly trained with sufficient data, the neural ODE
vector field ${\bf f}(x,t;\theta_{k})$ will learn to approximate (12). This
can be viewed as inferring the score function of $\rho_{t}$, and also leads to
the invertibility of the trained flow net in theory: Suppose the trained ${\bf
f}(x,t;\theta_{k})$ is close enough to (12), it will also has bounded
Lipschitz constant. Then the residual block is invertible as long as the step
size $h$ is sufficiently small, e.g. less than $1/L$ where $L$ is the
Lipschitz bound of ${\bf f}(x,t;\theta_{k})$. In practice, we typically use
smaller $h$ than needed merely by invertibility (allowed by model budget) so
that the flow network can more closely track the Fokker-Planck equation of the
diffusion process. The invertibility of the proposed model is numerically
verified in experiments (see Table 1).

##  4 Training of JKO-iFlow net

The proposed JKO-iFlow model allows progressive learning of the residual
blocks in the neural-ODE model in a block-wise manner (Section 4.1). We also
introduce two techniques to improve the training of the trajectories in
probability space (Section 4.2), which can be illustrated in a gradient system
in vector space (Appendix B.2).

###  4.1 Block-wise training

Note that the training of $(k+1)$-th block in (8) can be conducted once the
previous $k$ blocks are trained. Specifically, with finite training data
$\\{X_{i}=x_{i}(0)\\}_{i=1}^{n}$, the expectation $\mathbb{E}_{x(t)\sim
p_{k}}$ in (8) is replaced by the sample average over
$\\{x_{i}(kh)\\}_{i=1}^{n}$ which can be computed from the previous $k$
blocks. Note that for each given $x(t)=x(t_{k})$, both $x(t_{k+1})$ and the
integral of $\nabla\cdot{\bf f}$ in (8) can be computed by a numerical neural
ODE integrator. Following previous works, we use the Hutchinson trace
estimator (Hutchinson, 1989; Grathwohl et al., 2019) to estimate the quantity
$\nabla\cdot{\bf f}$ in high dimensions. Applying the numerical integrator in
computing (8), we denote the resulting $k$-th residual block abstractly as
$f_{\theta_{k}}$ with trainable parameters $\theta_{k}$.

Algorithm 1 Block-wise JKO-iFlow training

0: Time stamps $\\{t_{k}\\}$, training data, termination criterion $\rm{Ter}$
and tolerance level $\epsilon$, maximal number of blocks $L_{\max}$.

1: Initialize $k=1$.

2: while $\rm{Ter}(k)>\epsilon$ and $k\leq L_{\max}$ do

3: Optimize $f_{\theta_{k}}$ upon minimizing (8) with mini-batch sample
approximation, given $\\{f_{\theta_{i}}\\}_{i=1}^{k-1}$. Set $k\leftarrow
k+1$.

4: end while

5: $L\leftarrow k$. Optimize $f_{\theta_{L+1}}$ using (8) without the $W_{2}$
regularization. {$\triangleright$ Free block.}

This leads to a block-wise training of the normalizing flow network, as
summarized in Algorithm 1. The sequence of time stamps $t_{k}$ is given by
specifying the time steps $h_{k}:=t_{k+1}-t_{k}$, which we allow to differ
across $k$. The choice of the sequence $h_{k}$ is initialized by a geometric
sequence starting from $h_{0}$ with maximum stepsize $h_{\rm max}$, cf.
Appendix B.1. In the special case where the multiplying factor is one, the
sequence of $h_{k}$ gives a constant step size. The adaptive choice of $h_{k}$
with refinement (by adding more blocks) is to be introduced in Section 4.2.

Regarding the termination criterion $\rm{Ter}(k)$ in line 2 of Algorithm 1, we
monitor the ratio $\mathbb{E}_{x\sim
p_{k-1}}\|x-T_{k}(x)\|^{2}/\mathbb{E}_{x\sim p_{k-1}}\|T_{k}(x)\|^{2}$ and
terminate when it is below some threshold $\epsilon$, set as 0.01 in all
experiments. Lastly, line 5 of training a “free block” (i.e., block without
the $W_{2}$ regularization) is to flow the push-forward density $p_{L}$ closer
to the target density $p_{Z}$, where the former is obtained through the first
$L$ blocks.

The block-wise training significantly reduces memory and computational
complexity, since only one block is trained when optimizing (8) regardless of
flow depth. Therefore, one can use larger training batches and potentially
more expensive numerical integrators within a certain memory budget for higher
accuracy. We also empirically observe that training each block using standard
back-propagation (rather than the adjoint method in neural ODE) gives a
comparable result at a lower cost. At last, to ensure the the invertibility of
the trained JKO-iFlow network, we furtherly break the time interval
$[t_{k-1},t_{k})$ into 3 or 5 subintervals to compute the neural ODE
integration e.g. by Runge-Kutta-4. We empirically verify the smallness of
inversion errors on test samples in Section 5.

###  4.2 Improved computation of trajectories in probability space

We adopt two additional computational techniques to facilitate learning of the
trajectories in the probability space, represented by the sequence of
densities $p_{k}$, $k=1,\cdots,L$, associated with the $L$ residual blocks of
the proposed normalizing flow network. The two techniques are illustrated in
Figure 2. Further details and illustrations of the approach in vector space
can be found in Appendix B.

$\bullet$ Trajectory reparameterization. We empirically observe fast decay of
the movements $W_{2}^{2}(T_{\\#}p_{k},p_{k})$ when $h_{k}$ is set to be
constant, that is, initial blocks transport the densities much further than
the later ones. This is consistent with the exponential convergence of the
Fokker-Planck flow, cf. (4), but unwanted in algorithm because in order to
train the current block, the flow model needs to transport data through all
previous blocks, and thus the later blocks trained using constant step size
barely contribute to the density transport. Hence, instead of having constant
$h_{k}$, we reparameterize the values of $t_{k}$ through an adaptive procedure
based on the $W_{2}$ distance at each block. The adaptive procedure uses an
interactive approach to encourage the W2 movement in each block to be more
even across the $L$ blocks, where the retraining of the trajectory can be
potentially warm-started by the previous trajectory in iteration.

$\bullet$ Progressive refinement. The performance of CNF models is typically
improved with larger number of residual blocks, corresponding to smaller step
size $h_{k}$. The smallness of stepsize $h$ also ensures the invertibility of
the flow model, in theory and also in practice. However, directly training the
model with a small non-adaptive stepsize $h$ may result in long computation
time and convergence to the normal density $q_{Z}$ only after a large number
of blocks, where the choice of $h_{k}$ is not as efficient as after adaptive
reparameterization. We introduce a refinement approach to increase the number
of blocks progressively, where each time interval $[t_{k-1},t_{k})$ splits
into two halves and the number of blocks doubles after the adaptive
reparameterization of the trajectory converges at the coarse level. The new
trajectory at the refined level is again trained with adaptive
reparameterization, where the residual blocks can be warm-started from the
coarse-level model to accelerate the convergence. The trajectory refinement
allows going to a smaller step size $h_{k}$, which benefits the accuracy of
the JKO-iFlow model including the numerical accuracy in integrating each
neural ODE block.

Figure 2:  Diagram illustrating trajectory reparameterization and refinement.
The top panel shows the original trajectory under three blocks via Algorithm
1. The bottom panel shows the trajectory under six blocks after
reparameterization and refinement, which renders the W2 movements more even.

##  5 Experiment

In this section, we examine the proposed JKO-iFlow model on simulated and real
datasets, including both unconditional and conditional generation tasks. Codes
will be released upon publication.

a True data JKO-iFlow
$\tau$: 3.10e-4, MMD-c: 3.09e-4
NLL 2.64

b FFJORD
4.47e-4
2.95

c OT-Flow
1.62e-03
3.30

d $W_{2}$ IResNet
3.34e-3
3.35

e ScoreSDE
7.94e-4
3.2

f Fractal tree
$\tau$: 3.12e-4, MMD-c: 2.17e-4
NLL 2.20

g Olympic rings
$\tau$: 3.16e-4, MMD-c: 2.36e-4
NLL 1.66

h Checkerboard
$\tau$: 3.09e-4, MMD-c: 2.70e-4
NLL 3.59

Figure 3:  Results on two-dimensional simulated datasets by JKO-iFlow and
competitors.  Table 1:  Inversion error $\mathbb{E}_{x\sim
p_{X}}\|{T_{\theta}^{-1}}({T_{\theta}}(x))-x\|^{2}_{2}$ of JKO-iFlow computed
via sample average on the test split of the data set, where $T_{\theta}$
denotes the transport mapping over all the blocks of the trained flow network.

Rose | Fractal tree | Olympic rings | Checkerboard | POWER | GAS | MINIBOONE | BSD300 | MNIST
---|---|---|---|---|---|---|---|---
3.30e-6 | 3.58e-5 | 2.24e-6 | 3.07e-5 | 1.48e-5 | 1.58e-6 | 1.09e-6 | 1.53e-5 | 1.87e-5

###  5.1 Baselines and metrics

We compare with five alternatives, including four CNF models and one diffusion
model. The first two CNF models are FFJORD (Grathwohl et al., 2019) and OT-
Flow (Onken et al., 2021), which are continuous-time flow (neural-ODE models).
The next two CNF models are IResNet (Behrmann et al., 2019) and IGNN (Xu et
al., 2022), which are discrete in time (ResNet models). We call the model in
IGNN the ‘$W_{2}$ IResNet’ as it applies Wasserstein-2 regularization in
IResNet. The diffusion model baseline is the score-based neural SDE (Song et
al., 2021), which we call ‘ScoreSDE’. Detailed information about experimental
setup, including dataset construction and neural network architecture and
training, for each dataset can be found in Appendix C.2.

The accuracy of trained generative models is evaluated by two quantitative
metrics, the negative log-likelihood (NLL) metric, and the kernel maximum mean
discrepancy (MMD) (Gretton et al., 2012) metric, including MMD-1, MMD-m, and
MMD-c for constant, median distance, and custom bandwidth respectively. The
test threshold $\tau$ is computed by bootstrap, where an MMD metric less than
$\tau$ indicates that the generated distribution is evaluated by the MMD test
to be the same as the true data distribution (achieving $\alpha=0.05$ test
level). See details in Appendix C.1. The computational cost is measured by the
number of mini-batch stochastic gradient descent steps (e.g., batches that the
model uses in training) and the training time. The performance is also
reported under a fixed-budget setting to compare across models. At last, we
track the numerical inversion error of the trained flow model on all the
datasets, cf. Table 1.

Table 2:  MMD metrics on high-dimensional real datasets, where all competitors
are trained in a fixed-budget setup using 10 times more mini-batches (their
performances using the same number of mini-batches are worse and not
comparable to JKO-iFlow). Complete results, including NLL metric and baselines
of competitors using more computation, are given in Table A.1.

Data Set | Model | # Param | Test MMD-m | Test MMD-1
---|---|---|---|---
POWER $d=6$ |  |  |  $\tau$: 8.33e-4 |  $\tau$: 1.43e-3
JKO-iFlow | 76K | 8.21e-4 | 1.26e-3
OT-Flow | 76K | 5.69e-4 | 9.62e-4
FFJORD | 76K | 1.38e-3 | 1.98e-3
$W_{2}$ IResNet | 304K | 2.76e-3 | 2.72e-3
IResNet | 304K | 4.50e-3 | 2.49e-2
ScoreSDE | 76K | 1.34e-3 | 6.56e-3
GAS $d=8$ |  |  |  $\tau$: 9.64e-4 |  $\tau$: 1.46e-3
JKO-iFlow | 76K | 5.96e-4 | 1.79e-3
OT-Flow | 76K | 1.51e-3 | 3.64e-3
FFJORD | 76K | 3.62e-3 | 6.09e-3
$W_{2}$ IResNet | 304K | 7.14e-3 | 1.50e-2
IResNet | 304K | 3.26e-3 | 2.72e-2
ScoreSDE | 76K | 1.31e-3 | 1.45e-3

Data Set | Model | # Param | Test MMD-m | Test MMD-1
---|---|---|---|---
MINIBOONE $d=43$ |  |  |  $\tau$: 7.76e-4 |  $\tau$: 1.11e-3
JKO-iFlow | 112K | 7.97e-4 | 1.01e-3
OT-Flow | 112K | 1.23e-3 | 1.01e-3
FFJORD | 112K | 5.47e-3 | 1.04e-3
$W_{2}$ IResNet | 448K | 1.27e-2 | 1.03e-3
IResNet | 448K | 2.58e-3 | 1.04e-3
ScoreSDE | 112K | 4.29e-3 | 1.03e-3
BSDS300 $d=63$ |  |  |  $\tau$: 6.72e-4 |  $\tau$: 4.55e-4
JKO-iFlow | 396K | 4.83e-3 | 3.03e-3
OT-Flow | 396K | 8.55e-2 | 8.44e-2
FFJORD | - | - | -
$W_{2}$ IResNet | 990K | 5.52e-1 | 6.88e-1
IResNet | 990K | 5.42e-1 | 5.95e-1
| ScoreSDE | 396K | 5.51e-1 | 6.62e-1

###  5.2 Two-dimensional toy data

The results on four toy datasets are shown in Figure 3, where the metrics of
NLL and MMD-c are reported. Plots (a)-(e) compare JKO-iFlow with the
alternative models on the dataset Rose, where both NLL and MMD-c metrics show
the better performance of JKO-iFlow. Visually, the generated samples $\hat{X}$
by JKO-iFlow also more resemble the true data distribution of $X$ than the
competitors in (b)-(e). Plots (f)-(h) show the results by JKO-iFlow on the
examples of Fractal tree, Olympic rings, and Checkerboard, where the MMD-c
metrics by JKO-iFlow are all below the threshold $\tau$, indicating
satisfactory generative performance.

The reported performance by JKO-iFlow in Figure 3 is obtained after applying
the trajectory improvement techniques in Section 4.2. The performance before
applying reparameterization and refinement is provided in Figures A.2 and A.3,
which shows the benefit of trajectory improvement in learning the model.

###  5.3 High-dimensional tabular data

Table 2 assesses the performance of JKO-iFlow and baseline methods on four
high-dimensional real datasets under a fixed-budget setting. (Additional
baselines of OT-Flow and FFJORD from the original references are included in
Table A.1). For a fair comparison, we keep the model size of continuous flow
models (OT-Flow and FFJORD) and the diffusion model (ScoreSDE) the same, and
we increase model sizes for discrete flow models ($W_{2}$ IResNet and IResNet)
due to that each residual block can be computed faster and the testing
performance can be improved using larger models. The FFJORD baseline on
BSDS300 is omitted due to a significantly longer training time before model
convergence.

In terms of performance, except for POWER dataset, JKO-iFlow yields smaller or
similar MMD-m and MMD-1 metrics compared to all alternatives. The visual
comparisons of generated data distributions (projected to 2D) are shown in
Figure A.4, where we note that the evaluation by MMD-m more aligns with the
visual performance than the MMD-1 metric. The results show that JKO-iFlow
gives the overall more superior performance on the tabular datasets.

Furthermore, the results of applying the trajectory reparameterization on
MINIBOONE dataset are shown in Figure 4. The benefit appears in yielding a
flow trajectory with more uniform $W_{2}$ movement, as shown in plot (a). The
generative performances before and after the reparameterization iterations are
comparable, which is shown in plot (b), and the MMD-m metrics in both cases
are under the threshold.

a Per-block $W_{2}^{2}$ over reparameterization iterations.

b Results at initial training (middle) and Iter 4 (right).

Figure 4:  Reparametrization iterations of JKO-iFlow model on MINIBOONE
dataset. In (a), results labeled “Iter 4” are obtained after 4
reparameterization iterations, which yield more uniform $W_{2}$ movements
compared to the initial trajectory labeled “Initial”. The generated samples by
the models before and after the moving iterations are shown in (b).

###  5.4 MNIST

We apply the JKO-iFlow model to MNIST data after mapping the input images to a
code space of dimension 16 using a pre-trained Auto Encoder (AE). The pre-
trained AE gives the encoder ${\rm
Enc}:\mathbb{R}^{1024}\rightarrow\mathbb{R}^{d}$ and decoder ${\rm
Dec}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{1024}$ such that ${\rm Dec}({\rm
Enc}(X))\approx X$ for $X$ re-sized to 32-by-32 pixels, where we choose
$d=16$. Then, JKO-iFlow is trained on $N$=60k encoded images $\\{{\rm
Enc}(X_{i})\\}_{i=1}^{N}$ to give an invertible transport mapping (over all
residual blocks) $T_{\theta}:\mathbb{R}^{d}\to\mathbb{R}^{d}$. The images are
generated upon sampling noises $Z\sim\mathcal{N}(0,I_{d})$ through the
backward flow followed by the decoder, namely ${\rm
Dec}({T^{-1}_{\theta}}(Z))$. The generated images are shown in Figure 5. The
test MMD-m metric is 2.33e-4, versus the threshold $\tau$=2.59e-4.

Figure 5:  Generated MNIST samples by JKO-iFlow trained using a pre-trained
auto-encoder.

###  5.5 Conditional generation

The conditional generation problem aims to generate input samples $X$ given
labels $Y$ from the unknown conditional distribution $X|Y$ to be learned from
data. We follow the IGNN approach (Xu et al., 2022) and apply the JKO-iFlow to
toy data sets (Figure A.5) and the Solar ramping dataset which is on a graph
(Figure 6). The conditional generative CNF model considers categorical labels
$Y$ taking value in $K$ classes, and trains an invertible flow model to
transport the conditional data density $p_{X|Y=k}$ to $p_{H|Y=k}$ for all
class label $k$, where $H|Y$ is a Gaussian mixture model (GMM) with possibly
trainable parameters. To use JKO-iFlow for conditional generation in this
setting, we modify the objective (8) to use the GMM target distribution
instead of a standard normal $q_{Z}$. Specifically, instead of using
$V_{Z}(X)=\|X\|^{2}_{2}$, we replace $V_{Z}(X)$ with $V_{H|Y}(X|Y)=\log
p_{H|Y}(X|Y)$, where $H|Y\sim H(\mu_{Y},\sigma^{2}I_{d})$.

We compare JKO-iFlow with IGNN on the Solar ramping data, which is a graph
node prediction dataset, where bot model use graph neural networks in the
residual blocks. Figure 6 shows that JKO-iFlow yields better NLL and MMD
metrics than IGNN. The MMD values for both methods are above the threshold
$\tau$, and thus the smaller the MMD metric the better the performance.

a True $X|Y$ JKO-iFlow
$\tau$: 2.78e-3, MMD-m: 1.95e-2
NLL -0.33

b IGNN
2.60e-2
2.46

c True $X|Y$ JKO-iFlow
$\tau$: 1.30e-2, MMD-m: 5.89e-2
NLL -8.86

d IGNN
1.09e-1
-3.80

Figure 6:  Conditional graph node feature generation by JKO-iFlow and iGNN,
where nodal feature $X\in\mathbb{R}^{10\times 2}$. We visualize the
conditionally generated samples upon projecting down to the first two
principal components determined by true samples from $X|Y$. We visualize
generation at two different values of $Y$ denoting nodal labels.

##  6 Discussion

The work can be extended in several directions. The application to larger-
scale image datasets by adopting convolutional layers will further verify the
usefulness of the proposed method. The applications to generative tasks on
graph data, by incorporating graph neural network layers in JKO-iFlow model,
are also of interest. This also includes conditional generative tasks, of
which the first results on toy data are shown in this work. For the
methodology, the time continuity over the parametrization of the residual
blocks (as a result of the smoothness of the Fokker-Planck flow) has not been
exploited in this work, which may further improve model capacity as well as
learning efficiency.

Theoretically, the model expressiveness of the flow model to generate any
regular data distribution can be analyzed based on Section 3.3. To sketch a
road map, a block-wise approximation guarantee of ${\bf f}(x,t)$ as in (12)
can lead to approximation of the Fokker-Planck flow (3), which pushes forward
the density to be $\epsilon$-close to normality in $T=\log(1/\epsilon)$ time,
cf. (4). Reversing the time of the ODE then leads to an approximation of the
initial density $\rho_{0}=p_{X}$ by flowing backward in time from $T$ to zero.
Further analysis under technical assumptions is left to future work.

## Acknowledgement

The work is partially supported by NSF DMS-2134037. C.X. and Y.X. are
partially supported by an NSF CAREER CCF-1650913, and NSF DMS-2134037,
CMMI-2015787, CMMI-2112533, DMS-1938106, and DMS-1830210. X.C. is partially
supported by NSF and the Alfred P. Sloan Foundation.

## References

  * Alvarez-Melis et al. (2021) David Alvarez-Melis, Yair Schiff, and Youssef Mroueh.  Optimizing functionals on the space of probabilities with input convex neural networks.  _arXiv preprint arXiv:2106.00774_ , 2021.
  * Amos et al. (2017) Brandon Amos, Lei Xu, and J Zico Kolter.  Input convex neural networks.  In _International Conference on Machine Learning_ , pages 146–155. PMLR, 2017.
  * Arcones and Gine (1992) Miguel A Arcones and Evarist Gine.  On the bootstrap of u and v statistics.  _The Annals of Statistics_ , pages 655–674, 1992.
  * Bailey and Telgarsky (2018) Bolton Bailey and Matus J Telgarsky.  Size-noise tradeoffs in generative networks.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Barron (1993) Andrew R Barron.  Universal approximation bounds for superpositions of a sigmoidal function.  _IEEE Transactions on Information theory_ , 39(3):930–945, 1993.
  * Behrmann et al. (2019) Jens Behrmann, Will Grathwohl, Ricky TQ Chen, David Duvenaud, and Jörn-Henrik Jacobsen.  Invertible residual networks.  In _International Conference on Machine Learning_ , pages 573–582. PMLR, 2019.
  * Block et al. (2020) Adam Block, Youssef Mroueh, and Alexander Rakhlin.  Generative modeling with denoising auto-encoders and langevin sampling.  _arXiv preprint arXiv:2002.00107_ , 2020.
  * Bolley et al. (2012) François Bolley, Ivan Gentil, and Arnaud Guillin.  Convergence to equilibrium in wasserstein distance for fokker–planck equations.  _Journal of Functional Analysis_ , 263(8):2430–2457, 2012.
  * Bunne et al. (2022) Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi.  Proximal optimal transport modeling of population dynamics.  In _International Conference on Artificial Intelligence and Statistics_ , pages 6511–6528. PMLR, 2022.
  * Chen et al. (2018) Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.  Neural ordinary differential equations.  _Advances in neural information processing systems_ , 31, 2018.
  * Chen et al. (2019) Ricky TQ Chen, Jens Behrmann, David K Duvenaud, and Jörn-Henrik Jacobsen.  Residual flows for invertible generative modeling.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Defferrard et al. (2016) Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst.  Convolutional neural networks on graphs with fast localized spectral filtering.  _Advances in neural information processing systems_ , 29, 2016.
  * Dinh et al. (2014) Laurent Dinh, David Krueger, and Yoshua Bengio.  Nice: Non-linear independent components estimation.  _arXiv preprint arXiv:1410.8516_ , 2014.
  * Dinh et al. (2017) Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real nvp.  _ArXiv_ , abs/1605.08803, 2017.
  * Fan et al. (2021) Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen.  Variational wasserstein gradient flow.  _arXiv preprint arXiv:2112.02424_ , 2021.
  * Fey and Lenssen (2019) Matthias Fey and Jan E. Lenssen.  Fast graph representation learning with PyTorch Geometric.  In _ICLR Workshop on Representation Learning on Graphs and Manifolds_ , 2019.
  * Finlay et al. (2020) Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman.  How to train your neural ode: the world of jacobian and kinetic regularization.  In _International conference on machine learning_ , pages 3154–3164. PMLR, 2020.
  * Goodfellow et al. (2014) Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio.  Generative adversarial nets.  In _NIPS_ , 2014.
  * Grathwohl et al. (2019) Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Kristjanson Duvenaud.  Ffjord: Free-form continuous dynamics for scalable reversible generative models.  _ArXiv_ , abs/1810.01367, 2019.
  * Gretton et al. (2012) Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alex Smola.  A kernel two-sample test.  _J. Mach. Learn. Res._ , 13:723–773, 2012.
  * Gulrajani et al. (2017) Ishaan Gulrajani, Faruk Ahmed, Martín Arjovsky, Vincent Dumoulin, and Aaron C. Courville.  Improved training of wasserstein gans.  In _NIPS_ , 2017.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Huang et al. (2021) Chin-Wei Huang, Jae Hyun Lim, and Aaron C Courville.  A variational perspective on diffusion-based generative models and score matching.  _Advances in Neural Information Processing Systems_ , 34:22863–22876, 2021.
  * Hutchinson (1989) Michael F. Hutchinson.  A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines.  _Communications in Statistics - Simulation and Computation_ , 18:1059–1076, 1989.
  * Isola et al. (2017) Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A. Efros.  Image-to-image translation with conditional adversarial networks.  _2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 5967–5976, 2017.
  * Johnson and Zhang (2019) Rie Johnson and Tong Zhang.  A framework of composite functional gradient methods for generative adversarial models.  _IEEE transactions on pattern analysis and machine intelligence_ , 43(1):17–32, 2019.
  * Jordan et al. (1998) Richard Jordan, David Kinderlehrer, and Felix Otto.  The variational formulation of the fokker–planck equation.  _SIAM journal on mathematical analysis_ , 29(1):1–17, 1998.
  * Kingma and Ba (2015) Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _CoRR_ , abs/1412.6980, 2015.
  * Kingma and Welling (2014) Diederik P. Kingma and Max Welling.  Auto-encoding variational bayes.  _CoRR_ , abs/1312.6114, 2014.
  * Kingma and Welling (2019) Diederik P. Kingma and Max Welling.  An introduction to variational autoencoders.  _Foundations and Trends® in Machine Learning_ , 12(4):307–392, 2019.  doi: 10.1561/2200000056.
  * Kingma and Dhariwal (2018) Durk P Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  _Advances in neural information processing systems_ , 31, 2018.
  * Kobyzev et al. (2020) Ivan Kobyzev, Simon JD Prince, and Marcus A Brubaker.  Normalizing flows: An introduction and review of current methods.  _IEEE transactions on pattern analysis and machine intelligence_ , 43(11):3964–3979, 2020.
  * Korotin et al. (2021) Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev.  Wasserstein-2 generative networks.  In _International Conference on Learning Representations_ , 2021.
  * Lee et al. (2017) Holden Lee, Rong Ge, Tengyu Ma, Andrej Risteski, and Sanjeev Arora.  On the ability of neural nets to express distributions.  In _Conference on Learning Theory_ , pages 1271–1296. PMLR, 2017\.
  * Liu et al. (2019) Jenny Liu, Aviral Kumar, Jimmy Ba, Jamie Kiros, and Kevin Swersky.  Graph normalizing flows.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Liu and Wang (2016) Qiang Liu and Dilin Wang.  Stein variational gradient descent: A general purpose bayesian inference algorithm.  _Advances in neural information processing systems_ , 29, 2016.
  * Liutkus et al. (2019) Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Robert Stöter.  Sliced-wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions.  In _International Conference on Machine Learning_ , pages 4104–4113. PMLR, 2019.
  * Lu and Lu (2020) Yulong Lu and Jianfeng Lu.  A universal approximation theorem of deep neural networks for expressing probability distributions.  _Advances in neural information processing systems_ , 33:3094–3105, 2020.
  * Lucas et al. (2019) James Lucas, G. Tucker, Roger B. Grosse, and Mohammad Norouzi.  Understanding posterior collapse in generative latent variable models.  In _DGS@ICLR_ , 2019.
  * Mathieu and Nickel (2020) Emile Mathieu and Maximilian Nickel.  Riemannian continuous normalizing flows.  _Advances in Neural Information Processing Systems_ , 33:2503–2515, 2020.
  * Mokrov et al. (2021) Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Burnaev.  Large-scale wasserstein gradient flows.  _Advances in Neural Information Processing Systems_ , 34:15243–15256, 2021.
  * Onken et al. (2021) Derek Onken, S 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, 2021.
  * Papamakarios et al. (2017) George Papamakarios, Theo Pavlakou, and Iain Murray.  Masked autoregressive flow for density estimation.  _Advances in neural information processing systems_ , 30, 2017.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, _Advances in Neural Information Processing Systems 32_ , pages 8024–8035. Curran Associates, Inc., 2019.
  * Perekrestenko et al. (2020) Dmytro Perekrestenko, Stephan Müller, and Helmut Bölcskei.  Constructive universal high-dimensional distribution generation through deep relu networks.  In _International Conference on Machine Learning_ , pages 7610–7619. PMLR, 2020.
  * Perekrestenko et al. (2021) Dmytro Perekrestenko, Léandre Eberhard, and Helmut Bölcskei.  High-dimensional distribution generation through deep neural networks.  _Partial Differential Equations and Applications_ , 2(5):1–44, 2021.
  * Rout et al. (2022) Litu Rout, Alexander Korotin, and Evgeny Burnaev.  Generative modeling with optimal transport maps.  In _International Conference on Learning Representations_ , 2022.
  * Salimans et al. (2016) Tim Salimans, Ian J. Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen.  Improved techniques for training gans.  _ArXiv_ , abs/1606.03498, 2016.
  * Sideris (2013) Thomas C. Sideris.  Ordinary differential equations and dynamical systems.  2013\.
  * Song and Ermon (2019) Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Song et al. (2021) Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  In _International Conference on Learning Representations_ , 2021.
  * Stein (1972) Charles Stein.  A bound for the error in the normal approximation to the distribution of a sum of dependent random variables.  In _Proceedings of the sixth Berkeley symposium on mathematical statistics and probability, volume 2: Probability theory_ , volume 6, pages 583–603. University of California Press, 1972.
  * Tabak and Vanden-Eijnden (2010) Esteban G Tabak and Eric Vanden-Eijnden.  Density estimation by dual ascent of the log-likelihood.  _Communications in Mathematical Sciences_ , 8(1):217–233, 2010.
  * Tzen and Raginsky (2019a) Belinda Tzen and Maxim Raginsky.  Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit.  _arXiv preprint arXiv:1905.09883_ , 2019a.
  * Tzen and Raginsky (2019b) Belinda Tzen and Maxim Raginsky.  Theoretical guarantees for sampling and inference in generative models with latent diffusions.  In _Conference on Learning Theory_ , pages 3084–3114. PMLR, 2019b.
  * Xu et al. (2022) Chen Xu, Xiuyuan Cheng, and Yao Xie.  Invertible neural networks for graph prediction.  _arXiv preprint arXiv:2206.01163_ , 2022.
  * Yarotsky (2017) Dmitry Yarotsky.  Error bounds for approximations with deep relu networks.  _Neural Networks_ , 94:103–114, 2017.
  * Zhang et al. (2018) Linfeng Zhang, Lei Wang, et al.  Monge-amp$\backslash$ere flow for generative modeling.  _arXiv preprint arXiv:1809.10188_ , 2018.
  * Zhang and Chen (2021) Qinsheng Zhang and Yongxin Chen.  Diffusion normalizing flow.  _Advances in Neural Information Processing Systems_ , 34:16280–16291, 2021.

##  Appendix A Proofs

###  A.1 Proofs in Section

######  Lemma A.1.

Suppose $p$ and $q$ are two densities on $\mathbb{R}^{d}$ in $\mathcal{P}$,
the following two problems

$\min_{\rho\in\mathcal{P}}L_{\rho}[\rho]={\rm
KL}(\rho||q)+\frac{1}{2h}W_{2}^{2}(p,\rho),$ (13)
$\min_{T:\mathbb{R}^{d}\to\mathbb{R}^{d}}L_{T}[T]={\rm
KL}(T_{\\#}p||q)+\frac{1}{2h}\mathbb{E}_{x\sim p}\|x-T(x)\|^{2},$ (14)

have the same minimum, and

(a) If $T^{*}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is a minimizer of (14), then
$\rho^{*}=(T^{*})_{\\#}p$ is a minimizer of (13).

(b) If $\rho^{*}$ is a minimizer of (13), then the optimal transport from $p$
to $\rho^{*}$ minimizes (14).

###### Proof of Lemma A.1.

Let the minimum of (14) be $L_{T}^{*}$, and that of (13) be $L_{\rho}^{*}$.

Proof of (a): Suppose $L_{T}$ achieves minimum at $T^{*}$, then $T^{*}$ is the
optimal transport from $p$ to $\rho^{*}=(T^{*})_{\\#}p$ because otherwise
$L_{T}$ can be further improved. By definition of $L_{\rho}$, we have
$L_{T}^{*}=L_{T}[T^{*}]=L_{\rho}[\rho^{*}]\geq L_{\rho}^{*}$. We claim that
$L_{T}^{*}=L_{\rho}^{*}$. Otherwise, there is another $\rho^{\prime}$ such
that $L_{\rho}[\rho^{\prime}]<L_{T}^{*}$. Let $T^{\prime}$ be the optimal
transport from $p$ to $\rho^{\prime}$, and then
$L_{T}[T^{\prime}]=L_{\rho}[\rho^{\prime}]<L_{T}^{*}$, contradicting with that
$L_{T}^{*}$ is the minimum of $L_{T}$. This also shows that
$L_{\rho}[\rho^{*}]=L_{T}^{*}=L_{\rho}^{*}$, that is, $\rho^{*}$ is a
minimizer of $L_{\rho}$.

Proof of (b): Suppose $L_{\rho}$ achieves minimum at $\rho^{*}$. Let $T^{*}$
be the OT from $p$ to $\rho^{*}$, then $\mathbb{E}_{x\sim
p}|x-T^{*}(x)|^{2}=W_{2}(p,\rho^{*})^{2}$, and then
$L_{T}[T^{*}]=L_{\rho}[\rho^{*}]=L_{\rho}^{*}$ which equals $L_{T}^{*}$ as
proved in (a). This shows that $T^{*}$ is a minimizer of $L_{T}$. ∎

###### Proof of Proposition 3.1,.

Given $p_{k}$ being the density of $x(t)$ at $t=kh$, recall that $T$ is the
solution map from $x(t)$ to $x(t+h)$. We denote $\rho_{t}:=p_{k}$, and
$\rho_{t+h}:=T_{\\#}p_{k}$. By definition,

${\rm
KL}(T_{\\#}p_{k}||p_{Z})=\mathbb{E}_{x\sim\rho_{t+h}}(\log{\rho}_{t+h}(x)-\log
p_{Z}(x)).$ (15)

Because $p_{Z}\propto e^{-V}$, $V(x)={|x|^{2}}/{2}$, we have $\log
p_{Z}(x)=-V(x)+c_{1}$ for some constant $c_{1}$. Thus

$\mathbb{E}_{x\sim\rho_{t+h}}\log p_{Z}(x)=\mathbb{E}_{x(t)\sim\rho_{t}}\log
p_{Z}(x(t+h))=c_{1}-\mathbb{E}_{x(t)\sim\rho_{t}}V(x(t+h)).$ (16)

To compute the first term in (15), note that

$\mathbb{E}_{x\sim\rho_{t+h}}\log\rho_{t+h}(x)=\mathbb{E}_{x(t)\sim\rho_{t}}\log\rho_{t+h}(x(t+h)),$
(17)

and by the expression (called “instantaneous change-of-variable formula” in
normalizing flow literature [Chen et al., 2018], which we derive directly in
below)

${\frac{d}{dt}}\log\rho(x(t),t)=-\nabla\cdot{\bf f}(x(t),t),$ (18)

we have that for each value of $x(t)$,

$\displaystyle\log\rho_{t+h}(x(t+h))=\log\rho(x(t+h),t+h)$
$\displaystyle=\log\rho(x(t),t)-\int_{t}^{t+h}\nabla\cdot\mathbf{f}(x(s),s)ds.$

Inserting back to (17), we have

$\mathbb{E}_{x\sim\rho_{t+h}}\log\rho_{t+h}(x)=\mathbb{E}_{x(t)\sim\rho_{t}}\log\rho_{t}(x(t))-\mathbb{E}_{x(t)\sim\rho_{t}}\int_{t}^{t+h}\nabla\cdot\mathbf{f}(x(s),s)ds).$

The first term is determined by $\rho_{t}=p_{k}$, and thus is a constant
$c_{2}$ independent from ${\bf f}(x,t)$ on $t\in[kh,(k+1)h]$. Together with
(16), we have shown that

$\text{r.h.s. of \eqref{eq:KL-
pf-1}}=c_{2}-\mathbb{E}_{x(t)\sim\rho_{t}}\int_{t}^{t+h}\nabla\cdot\mathbf{f}(x(s),s)ds)-c_{1}+\mathbb{E}_{x(t)\sim\rho_{t}}V(x(t+h)),$

which proves (7).

Derivation of (18): by chain rule,

$\displaystyle{\frac{d}{dt}}\log\rho(x(t),t)$
$\displaystyle=\frac{\nabla\rho(x(t),t)\cdot\dot{x}(t)+\partial_{t}\rho(x(t),t)}{\rho(x(t),t)}$
$\displaystyle=\left.\frac{\nabla\rho\cdot{\bf f}-\nabla\cdot(\rho{\bf
f})}{\rho}\right|_{(x(t),t)}\quad\text{(by \eqref{eq:ode-Xt} and
\eqref{eq:Liouville})}$ $\displaystyle=-\nabla\cdot{\bf f}(x(t),t).$

∎

##  Appendix B Technical details of Section 4

To train the proposed JKO-iFlow model by Algorithm 1, one needs to specify a
sequence of $t_{k}$, where the $k$-th JKO block integrates from $t_{k-1}$ to
$t_{k}$, starting from $t_{0}=0$. Denote by $h_{k}:=t_{k+1}-t_{k}$ the step-
size for $k$-th step.

###  B.1 Initial choice of $t_{k}$ in Algorithm 1

We first initialize the choice by the following scheme

$h_{k}=\min\\{\rho^{k}h_{0},h_{\rm max}\\},\quad k=0,1,\cdots$ (19)

where the base stepsize $h_{0}$, the multiplying factor $\rho\geq 1$, and the
maximum stepsize $h_{\rm max}\geq h_{0}$ are three hyper-parameters to be
specified by the user. When $\rho=1$, the sequence of $h_{k}$ in (19)
corresponds to constant stepsize $h_{k}\equiv h_{0}$.

Applying Algorithm 1 with the initial choice of $t_{k}$ gives $L$ trained
residual blocks of the neural ODE model, where $L$ is to be determined by the
stopping criterion in Section 4.1. The model is further improved by the two
additional techniques in Section 4.2, which will adaptively specify the
stepsize $h_{k}$ and we explain in detail below.

###  B.2  Trajectory improvement illustrated in vector space

Section 4.2 introduces two techniques to improve the computation of a gradient
descent trajectory represented by discrete points. For illustrative purposes,
we introduce the two techniques when the gradient system is in a vector space
equipped with Euclidean metric. The JKO-iFlow training applies the same
technique in the gradient system equipped with the Wasserstein-2 metric.

Let the space be $\mathbb{R}^{d}$, and $F:\mathbb{R}^{d}\to\mathbb{R}$ be a
differential landscape. Given a sequence of positive stepsize $h_{k}$, one can
compute a sequence of representative points $x_{k}$ which discretizes the
gradient descent trajectory of minimizing $F$ (towards a local minimum
$x^{*}$). Specifically, we have

$x_{k+1}=\arg\min_{x}F(x)+\frac{1}{2h_{k}}\|x-x_{k}\|_{2}^{2},$ (20)

starting from some fixed $x_{0}$.

####  B.2.1 Trajectory reparameterization

Starting from an initial sequence of step-size $\\{h_{k}\\}_{k=1}^{L}$ as in
Section B.1, the reparameterization technique implements an iterative scheme
to update the sequence of $h_{k}$ adaptively. We call the $j$-th iteration
‘Iter-$j$’, and denote the sequence by $\\{h_{k}^{(j)}\\}_{k=1}^{L}$. The
corresponding solution points $x_{k}$ by solving (20) with $h_{k}^{(j)}$ are
denoted as $x_{k}^{(j)}$. In addition to the constant $h_{\max}$, we need
another algorithmic parameter $\eta\in(0,1)$, and we set $\eta=0.3$ for the
vector-space example in computation.

We always have $x_{0}^{(j)}=x_{0}$ for all ‘Iter-$j$’. In the $j$-th iteration
of the parametrization, we compute the following

  1. 1.

Compute the arclength as

$S_{k}^{(j)}:=\|x_{k-1}^{(j)}-x_{k}^{(j)}\|_{2},\quad k=1,\cdots,L.$ (21)

  2. 2.

Compute the average arclength

$\bar{S}:=\sum_{k=1}^{L}S_{k}^{(j)}/L.$

  3. 3.

Update the stepsize for $k=1,\ldots,L$ as

$h_{k}^{(j+1)}:={\min\\{h_{k}^{(j)}+\eta(\bar{S}h_{k}^{(j)}/S_{k}^{(j)}-h_{k}^{(j)}),h_{\max}\\}}$

  4. 4.

Solve (20) using $h_{k}^{(j+1)}$ to obtain $x_{k}^{(j+1)}$.

We terminate the iteration when the arclengths $\\{S_{k}^{(j)}\\}$ are
approximately equal. An illustration is given as Iter-12 in the upper panel of
Figure A.1. After the reparameterization iterations, we solve an additional
$x_{L+1}$ by minimizing $F(x)$ starting from $x_{L}$. This corresponds to
optimizing the “free-block” in Algorithm 1.

####  B.2.2  Progressive refinement

The scheme can be computed using a positive integer refinement factor. For
simplicity, we use factor 2 throughout this work.

Given a sequence of $\\{h_{k}\\}_{k=1}^{L}$ and the representitive points
$\\{x_{k}\\}_{k=1}^{L}$, the refinement scheme computes a refined trajectory
having $k=1,\cdots,2L$:

  1. 1.

Compute the refined stepsizes as

$\tilde{h}_{2k-1}=\tilde{h}_{2k}=h_{k}/2,\quad k=1,\dots,L$

  2. 2.

Compute the representative points $\tilde{x}_{k}$ by solving (20) using
$\tilde{h}_{k}$, possibly warm-starting by $\tilde{x}_{2k}={x}_{k}$ and
$\tilde{x}_{2k-1}=({x}_{k}+x_{k-1})/2$.

We then apply the trajectory reparameterization iterations as in Section B.2.1
to the refined trajectory till convergence, and the free-block ending point is
also recomputed. An illustration is given as r-Iter-10 in the upper panel of
Figure A.1.

###  B.3 Trajectory improvement in probability space

We apply the two techniques to solve the JKO step (8), which is the
counterpart of (20) as a gradient descent scheme with proximal steps. The
representative points $x_{k}$ are replaced with transported distributions
$p_{k}$ which form a sequence of points in $\mathcal{P}$, and the optimization
is over the parametrization of $p_{k}$ induced by the neural network flow
mapping consisting of the first $k$ residual blocks.

It remains to define the arclength in (21) to implement the two techniques.
Because we equip $\mathcal{P}$ with the Wasserstein-2 metric, we compute
(omitting the iteration index $j$ in the notation)

$S_{k}=W_{2}(p_{k-1},p_{k})=(\mathbb{E}_{x\sim
p_{k-1}}\|x-T_{k}(x)\|^{2})^{1/2},$ (22)

where the transport mapping
$T_{k}(x)=x+\int_{t_{k-1}}^{t_{k}}f_{\theta_{k}}(x(s),s)ds$ and can be
computed from the $k$-th block. In practice, the expectation
$\mathbb{E}_{x\sim p_{k-1}}$ in (22) is computed by a finite-sample average on
the training set.

At last, the optimal warm-start of $p_{k}$ in the refinement is implemented by
inheriting the parameters $\theta_{k}$ of the trained blocks.

Figure A.1: The upper panel shows the arc lengths and the mean and standard
deviation of arc lengths over 12 reparameterization iterations, one
refinement, and an additional 10 reparameterization iterations. The lower
panel visualizes the trajectory consisting of 17 solution points at
“r-Iter-10”, where the free point as the 17-th point computes the appproximate
minimum by $x_{17}=x_{L+1}=[-1.911,0.105]$.

##  Appendix C Experimental details

###  C.1 Quantitative evaluation metrics

Besides visual comparison, we adopt two quantitative metrics to evaluate the
performance of generative models, the negative log-likelihood (NLL) metric
Grathwohl et al. [2019] and the maximum mean discrepancy (MMD) [Gretton et
al., 2012] metric.

####  C.1.1 NLL metric

Our computation of the NLL metric follows the standard procedure for neural
ODE normalizing flow models Chen et al. [2018], Grathwohl et al. [2019], where
the evolution of density $\rho(x(t))$ can be computed via integrating
$\nabla\cdot{\bf f}(x(t),t)$ over time due to the instantaneous change-of-
variable formula (18).

Specifically, our model flows from $t_{0}=0$ to $t_{L+1}=T$, where $x(0)\sim
p_{0}$ the data distribution, and we train the flow model to make $x(T)$
follow a normal density $q$. The model density at data sample $x$ is expressed
as $\rho(x,0)=(T_{\theta}^{-1})_{\\#}q(x)$ where $T_{\theta}^{-1}$ is the
inverse model flow mapping from $x(T)$ to $x(0)$. Thus the model log-
likelihood can be expressed as

$\log\rho(x(0),0)=\log q(x(T),T)+\int_{0}^{T}\nabla\cdot{\bf f}(x(s),s)ds.$

In practice, our trained flow model has $L$ blocks where each block
$f_{\theta_{k}}$ represents ${\bf f}(x,t)$ on $[t_{k-1},t_{k})$ and is
parametrized by $\theta_{k}$. The log-likelihood at test sample $x$ is then
computed by

${\rm
LL}(x)=-\frac{1}{2}(\|x(t_{L+1})\|^{2}_{2}+d\log(2\pi))+\sum_{k=1}^{L+1}\int_{t_{k-1}}^{t_{k}}\nabla\cdot
f_{\theta_{k}}(x(s),s)ds,$ (23)

where

$x(t_{k})=x(t_{k-1})+\int_{t_{k-1}}^{t_{k}}f_{\theta_{k}}(x(s),s)ds$

starting from $x(0)=x$. Both the integration of $f_{\theta_{k}}$ and
$\nabla\cdot f_{\theta_{k}}$ are computed using the numerical scheme of neural
ODE. We report NLL in the natural unit of information (i.e., $\log$ with base
$e$, known as “nats”) in all our experiments.

####  C.1.2 MMD metrics

Note that the normalizing flow models use NLL (on training set) as the
training objective. In contrast, the MMD metric is an impartial evaluation
metric as it is not used to train JKO-iFlow or any competing methods. Given
two set of data samples $\boldsymbol{X}:=\\{x_{i}\\}_{i=1}^{N}$ and
$\boldsymbol{Y}:=\\{y_{j}\\}_{j=1}^{M}$ and a kernel function $k(x,y)$, the
(squared) kernel MMD [Gretton et al., 2012] is defined as

$\text{MMD}(\boldsymbol{X},\boldsymbol{Y}):=\frac{1}{N^{2}}\sum_{i=1}^{N}\sum_{j=1}^{N}k(x_{i},x_{j})+\frac{1}{M^{2}}\sum_{i=1}^{M}\sum_{j=1}^{M}k(y_{i},y_{j})-\frac{2}{NM}\sum_{i=1}^{N}\sum_{j=1}^{M}k(x_{i},y_{j}),$
(24)

When a generative model is trained, we generate $M$ i.i.d. data samples by the
model to construct the set $\boldsymbol{Y}$, and we form the set
$\boldsymbol{X}$ using $N$ true data samples (from the test set). In all
experiments, we use the Gaussian kernel $k(x,y)=\exp\\{-\|x-y\|^{2}/2h^{2}\\}$
with bandwidth $h>0$. We use three ways of setting the bandwidth parameter
$h$:

  * •

Constant bandwidth: $h=h_{c}=1$. The resulting MMD is denoted as ‘MMD-1’.

  * •

Median bandwidth [Gretton et al., 2012]: let $h=h_{m}$ be the median of
$\|x_{i}-x_{j}\|$ for all distinct $i$ and $j$. The median distance is
computed from the dataset $X$. The resulting MMD is denoted as ‘MMD-m’.

  * •

Custom bandwidth: on certain datasets when we can use prior knowledge to
decide on the bandwidth, we will custom the choice of $h$ (typically smaller
than the median distance, due to that theoretically smaller bandwidth may lead
to a more powerful MMD test to distinguish the difference in the underlying
distributions) while ensuring that we use large enough $M$ and $N$ to compute
the MMD metric. We call the metric ‘MMD-c’.

Note that the MMD metric as a measure of distance between two distributions is
significant when above a test threshold $\tau$, which is defined as the upper
($1-\alpha$)-quantile of the distribution of the MMD statistic under the null
hypothesis (i.e., when dataset $\bf{X}$ and $\bf{Y}$ observe the same
distribution). The parameter $\alpha$ is the controlled Type-I error (known as
the test level), which is set to be 0.05. We compute the threshold $\tau$
using the bootstrap procedure following Arcones and Gine [1992], Gretton et
al. [2012], and we use 1000 times of bootstrap in all experiments. In our
usage of evaluating generative models, the threshold $\tau$ can be viewed a
baseline of the MMD metric, that is, when the computed MMD values are above
$\tau$, then the smaller the MMD value the better the generative performance;
when the computed MMD value is below $\tau$, it means that with respect to the
current MMD metric, the trained model generates a data distribution that is as
good as the true distribution.

###  C.2 Detail setup and additional results

All experiments are conducted using PyTorch [Paszke et al., 2019] and PyTorch
Geometric [Fey and Lenssen, 2019]. The optimizer is Adam [Kingma and Ba, 2015]
with a constant learning rate of 1e-3. We use the neural-ODE integrator
[Grathwohl et al., 2019, adjoint method] on two-dimensional examples and MNIST
with AutoEncoder, and use back-propagation on high-dimensional tabular
datasets.

a Per-block $W_{2}^{2}$ over reparameterization iterations.

b Results at initial training (middle) and Iter 3 (right). MMD and Negative
log-likelihood (NLL) are in title.

Figure A.2:  Same plots as in Figure 4 for Rose data. We observe improved
generative quality on the leaves of the rose after reparameterization
iterations

a Per-block $W_{2}^{2}$ over reparameterization iterations and refinement (cf.
refine and r-Iter 1, which indicates one reparameterization iteration after
refinement).

b Results at Iter 4 (middle) and r-Iter 1 (right). MMD and Negative log-
likelihood (NLL) are in title.

Figure A.3:  Same plots as in Figure 4 for Fractal tree data. We observe
improved generative quality on the edges of the tree after refinement.

####  C.2.1 Two-dimensional toy data

Training and test sets: We generate 16K test samples for each dataset, and for
the training split,

  * •

For Rose (Figure 3 and A.2) and Checkerboard and Olympic rings (Figure 3), we
re-sample 10K training samples every epoch, for a total of 100 epochs per
block. The batch size is 500.

  * •

For Fractal tree (Figure 3 and A.3), we use a fixed training data of 200K
samples, for a total of 50 epochs per block. The batch size is 2000.

Initial choice of $h_{k}$ schedule: $h_{0}=0.75,\rho=1.2$. We set $h_{\max}=5$
for rose, checkerboard, and Olympic rings, and $h_{\max}=3$ for fractal tree.

Reparametrization and refinement: we use $\eta=0.5$ in the reparameterization
iterations for all examples.

  * •

For Rose, 3 reparameterization moving iteration, no refinement.

  * •

For Checkerboard and Olympic rings, 4 reparameterization moving iteration, no
refinement.

  * •

For Fractal tree, 4 reparameterization moving iteration, one refinement, and
an additional reparameterization moving iteration.

MMD metric: We use custom kernel bandwidth $h=0.1h_{m}$ where $h_{m}$ is the
median distance of dataset $\bf{X}$ with $N=16000$, and we generate $M=N$ test
samples $\bf{Y}$. We then report the median of 100 MMD values, each of which
is computed by resampling without replacement $N/2$ samples from $\bf{X}$ and
$\bf{Y}$.

Network and activation: Fully-connected residual blocks with two hidden
layers. We use the softplus activation ($\beta=20$) with 128 hidden dimension.
Before refinement, we train 9 residual blocks for rose, checkerboard, and
Olympic rings, and train 6 residual blocks for fractal trees.

Additional results: Figure A.2 and A.3 illustrates the benefits of
reparameterization and refinement. Specifically, the details of the generated
rose around portions of leaves and of the generated tree around smallest edges
are clearer using these techniques. In terms of quantitative metrics, the NLL
is also smaller, even though the MMD metric indicates there is no statistical
difference between the ground truth and generated samples.

Table A.1:  Quantitative metrics on high-dimensional real datasets to augment
Table 2. The OT-Flow and FFJORD baselines marked with citations are from the
original paper (non-fixed budget setting), and the other baselines are
computed under the fixed budget setting.

Data Set | Model | # Param | Training | Testing
---|---|---|---|---
Time (h) | # Batches | Time/Batches (s) | Batch size | MMD-m | MMD-1 | NLL
POWER $d=6$ |  |  |  |  |  |  |  $\tau$: 8.33e-4 |  $\tau$: 1.43e-3 |
JKO-iFlow | 76K | 0.12 | 0.76K | 0.57 | 10000 | 8.21e-4 | 1.26e-3 | 0.58
OT-Flow | 76K | 1.21 | 7.58K | 0.57 | 10000 | 5.69e-4 | 9.62e-4 | 0.30
OT-Flow [Onken et al., 2021] | 18K | 3.1 | 22K | 0.56 | 10000 | - | 4.68e-5 | -0.30
FFJORD | 76K | 3.40 | 7.58K | 1.61 | 10000 | 1.38e-3 | 1.98e-3 | 0.60
FFJORD [Grathwohl et al., 2019] | 43K | - | - | - | 10000 | - | - | -0.46
$W_{2}$ IResNet | 304K | 0.49 | 7.58K | 0.23 | 10000 | 2.76e-3 | 2.72e-3 | 0.36
| IResNet | 304K | 0.76 | 7.58K | 0.36 | 10000 | 4.50e-3 | 2.49e-2 | 3.37
| ScoreSDE | 76K | 0.08 | 7.58K | 0.04 | 10000 | 1.34e-3 | 6.56e-3 | 3.41
| ScoreSDE | 76K | 0.84 | 75.85K | 0.04 | 10000 | 1.30e-3 | 5.66e-3 | 3.33
GAS $d=8$ |  |  |  |  |  |  |  $\tau$: 9.64e-4 |  $\tau$: 1.46e-3 |
JKO-iFlow | 76K | 0.08 | 0.76K | 0.38 | 5000 | 5.96e-4 | 1.79e-3 | -4.61
OT-Flow | 76K | 0.72 | 7.60K | 0.34 | 5000 | 1.51e-3 | 3.64e-3 | -4.29
OT-Flow [Onken et al., 2021] | 127K | 6.1 | 52K | 0.42 | 2000 | - | 2.47e-4 | -9.20
FFJORD | 76K | 3.49 | 7.60K | 1.65 | 5000 | 3.62e-3 | 6.09e-3 | -2.07
FFJORD [Grathwohl et al., 2019] | 279K | - | - | - | 1000 | - | - | -8.59
$W_{2}$ IResNet | 304K | 0.57 | 7.60K | 0.27 | 5000 | 7.14e-3 | 1.50e-2 | -4.45
| IResNet | 304K | 0.86 | 7.60K | 0.41 | 5000 | 3.26e-3 | 2.72e-2 | -1.17
| ScoreSDE | 76K | 0.04 | 7.60K | 0.02 | 5000 | 1.31e-3 | 1.45e-3 | -3.69
| ScoreSDE | 76K | 0.42 | 76.00K | 0.02 | 5000 | 4.27e-4 | 8.56e-4 | -5.58
MINIBOONE $d=43$ |  |  |  |  |  |  |  $\tau$: 7.76e-4 |  $\tau$: 1.11e-3 |
JKO-iFlow | 112K | 0.03 | 0.32K | 0.33 | 2000 | 7.97e-4 | 1.01e-3 | 13.63
OT-Flow | 112K | 0.75 | 3.39K | 0.80 | 2000 | 1.23e-3 | 1.01e-3 | 11.93
OT-Flow [Onken et al., 2021] | 78K | 0.8 | 7K | 0.44 | 2000 | - | 2.84e-4 | 10.55
FFJORD | 112K | 1.74 | 3.39K | 1.85 | 2000 | 5.47e-3 | 1.04e-3 | 23.45
FFJORD [Grathwohl et al., 2019] | 821K | - | - | - | 1000 | - | - | 10.43
$W_{2}$ IResNet | 448K | 0.80 | 3.25K | 0.89 | 2000 | 1.27e-2 | 1.03e-3 | 16.34
| IResNet | 448K | 1.32 | 3.25K | 1.46 | 2000 | 2.58e-3 | 1.04e-3 | 22.36
| ScoreSDE | 112K | 0.01 | 3.25K | 0.01 | 2000 | 4.29e-3 | 1.03e-3 | 27.38
| ScoreSDE | 112K | 0.09 | 32.48K | 0.01 | 2000 | 4.68e-3 | 1.10e-3 | 20.70
BSDS300 $d=63$ |  |  |  |  |  |  |  $\tau$: 6.72e-4 |  $\tau$: 4.55e-4 |
JKO-iFlow | 396K | 0.16 | 1.03K | 0.56 | 5000 | 4.83e-3 | 3.03e-3 | -156.67
OT-Flow | 396K | 3.50 | 10.29K | 1.22 | 1000 | 8.55e-2 | 8.44e-2 | -142.45
OT-Flow [Onken et al., 2021] | 297K | 7.1 | 37K | 0.70 | 300 | - | 4.24e-4 | -154.20
FFJORD | - | - | - | - | - | - | - | -
FFJORD [Grathwohl et al., 2019] | 6.7M | - | - | - | 10000 | - | - | -157.40
|  $W_{2}$ IResNet | 990K | 2.01 | 10.29K | 0.70 | 1000 | 5.52e-1 | 6.88e-1 | -107.39
| IResNet | 990K | 3.47 | 10.29K | 1.21 | 1000 | 5.42e-1 | 5.95e-1 | -33.11
| ScoreSDE | 396K | 0.01 | 10.29K | 0.005 | 1000 | 5.51e-1 | 6.62e-1 | -7.55
| ScoreSDE | 396K | 0.14 | 102.90K | 0.005 | 1000 | 5.51e-1 | 6.65e-1 | -7.31

a Power: JKO-iFlow

b OT-Flow

c FFJORD

d $W_{2}$ IResNet

e ScoreSDE

f Gas: JKO-iFlow

g OT-Flow

h FFJORD

i $W_{2}$ IResNet

j ScoreSDE

k MINIBOONE: JKO-iFlow

l OT-Flow

m FFJORD

n IResNet

o ScoreSDE

p BSDS300: JKO-iFlow

q OT-Flow

r $W_{2}$ IResNet

s IResNet

t ScoreSDE

Figure A.4: Generative quality on high-dimensional datasets via PCA projection
of generated samples. The generative quality in general aligns with the MMD-m
values shown in Table 2 and A.1.

####  C.2.2 Real datasets

Training and test sets: The four high-dimensional real datasets (POWER, GAS,
HEP- MASS, MINIBOONE) come from the University of California Irvine (UCI)
machine learning data repository, and we follow the pre-processing procedures
of [Papamakarios et al., 2017]. The solar dataset is retrieved from the
National Solar Radiation Database (NSRDB), following the setup as in [Xu et
al., 2022]. Regarding data sizes:

  * •

POWER: 1.85M training sample and 205K test sample.

  * •

GAS: 1M training sample, 100K test sample.

  * •

MINIBOONE: 32K training sample, 3.6K test sample.

  * •

BSDS300: 1.05M training sample, 250K test sample.

  * •

Solar: 1K training sample, 1K test sample.

Initial choice of $h_{k}$ schedule: $h_{0}=1,\rho=1,h_{\max}=5$.

Reparametrization and refinement: we use $\eta=0.5$ in the reparameterization
iterations for MINIBOONE.

MMD metric: We randomly draw 10% test samples for POWER ($N=20.5K$), GAS
($N=10K)$, and BSDS300 ($N=25K$), which are the same for JKO-iFlow and all
baseline methods. We use the entire test sample of MINIBOONE ($N=3.6K$) and
Solar ($N=1K$). We use the median distance and generate $M=N$ test sample. We
then report the median of 100 MMD values, each of which is computed by
resampling without replacement $N/2$ samples from $\bf{X}$ and $\bf{Y}$.

Network and activation: For the four high-dimensional tabular dataset, we use
fully-connected residual blocks with two hidden layers. On POWER, GAS, and
MINIBOONE, we use 128 hidden nodes, and on BSDS300, we use 256 hidden nodes.
For the solar data on graph, we design residual blocks to contain one Chebnet
[Defferrard et al., 2016] layer with degree 3, followed by two fully-connected
layers with 64 hidden nodes. We use the Tanh activation for all networks. For
results in Table 2, we train 4 residual blocks on all datasets, and for
results in Figure 4, we train 6 residual blocks.

Additional results: Figure A.4 visualizes the principal component projections
of the generated samples by JKO-iFlow and competitors of the high-dimensional
real datasets, where we often see a closer match between those from JKO-iFlow
with the ground truth. Table A.1 contains the complete numerical results of
JKO-iFlow against competitors on high-dimensional real datasets. For ScoreSDE,
we use the implementation in [Huang et al., 2021], which computes the evidence
lower bound (ELBO) for the data log-likelihood as reported in the last column.
Due to the block-wise training scheme, we measure the number of mini-batch
steps for JKO-iFlow as the sum of steps over all residual blocks that are
trained sequentially. Comparing to flow-based models, JKO-iFlow takes much
less iterations to reach a small enough MMD value. Although ScoreSDE is the
fastest, its performance, even under 100 times more mini-batch stochastic
gradient descent steps than JKO-iFlow, is still worse in terms of MMD-m on all
except GAS. We remark that JKO-iFlow yields larger NLL and MMD-1 than those in
previous works because we adopted fixed-budget computation (i.e., in terms of
the number of mini-batch gradient descent steps), rather than training the
residual blocks until convergence.

####  C.2.3 MNIST with pre-trained Auto-Encoder

Structure of AutoEncoder: The encoder/decoder consists of three
convolutional/convolutional transpose layers with filters $16\rightarrow
32\rightarrow 64$, under softplus activation before batch normalization. The
encoder/decoder also uses a fully-connected layer mapping to/from
$\mathbb{R}^{d}$ after/before convolutional/convolutional transpose layers.

Training and test sets: 60K training images and 10K test images. We train 500
epochs per block under a batch size of 2000.

Initial choice of $h_{k}$ schedule: $h_{0}=1,\rho=1,h_{\max}=3$.

Reparametrization and refinement: these techniques are not used.

Network and activation: Fully-connected residual blocks with two hidden
layers. We use the softplus activation ($\beta=20$) with 128 hidden dimension.
We train 6 residual blocks.

a Two-circles unconditional generation.

b Two-moons conditional generation.

Figure A.5:  Unconditional and conditional generation on simulated toy
datasets by JKO-iFlow. The setting of conditional generation problem is
introduced in Section 5.5.