Invertible normalizing flow neural networks by JKO scheme
Authors: Chen Xu
Xiuyuan Cheng
Yao Xie
Venue: Submitted to ICLR 2023 (2023)
Keywords: Normalizing flow
invertible neural networks
JKO scheme

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, past 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 an efficient \textit{block-wise} training procedure: as the JKO scheme unfolds the dynamic of gradient flow, the proposed model naturally stacks residual network blocks one-by-one and reduces the memory load as well as the difficulty of training deep networks. We also develop an adaptive time-reparametrization of the flow network with a progressive refinement of the trajectory in probability space, which improves the optimization efficiency and model accuracy in practice.  
On high-dimensional generative tasks for tabular data, JKO-Flow can process larger data batches and perform competitively as or better than continuous and discrete flow models, using 10X less number of iterations (e.g., batches) and significantly less time per iteration. 

Under review as a conference paper at ICLR 2023
INVERTIBLE NORMALIZING FLOW NEURAL NETWORKS BY JKO SCHEME
Anonymous authors Paper under double-blind review
ABSTRACT
Normalizing flow is 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, past works have regularized flow trajectories and designed special network architectures. The current paper develops a neural ODE flow network inspired by the Fordan-Kinderleehin allows an efficient block-wise training procedure: as the JKO scheme unfolds the dynamic of gradient flow, the proposed model naturally stacks residual network blocks one-by-one and reduces the memory load as well as the difficulty of training deep networks. We also develop an adaptive iime-reparamettiiation of the flow network with progressive refinement of the trajectory in probability space, which improves the optimization efficiency and model accuracy in practice. On high- dimensional generative tasks for tabular data, JKO-iFlow can process larger data batches and perform competitively as or better than continuous and discrete flow models, using 10X less number of iterations (e.g., batches) and significantly less time per iteration.
(JKO) scheme, which
INTRODUCTION
12% Jl0
Generative models have been widely studied in statistics and machine learning infer data-generating distributions and sample from the estimated distributions (Ronquist et al., 2012 Goodfellow et al. 2014)) Kingma & Welling, [014] Johnson & Zhang 2019). The normalizing flow has recently been a very popular generative framework. In short, a flow-based model learns the data distribution via an invertible mapping F between data density Px (X), 1, X E IRd and the target stan- dard multivariate Gaussian density px (Z))). Z (0, La) (Kobyzev et al. 2020). Benefits of the approach include effi- cient sampling and explicit likelihood computation. To make flow models practically useful, past works have made great efforts to develop flow models that facilitate training (e.g., in terms of loss objectives and computational techniques) and
induce smooth trajectories (Dinh et al. 2017 Grathwohl al. 2019 Onken et al. 2021).
223
(a) JKO-iFlow
(b) Generic Flow
Among flow models, continuous normalizing flow (CNF) trans- ports the data density to that of the target through continuous dynamics (e.g, Neural ODE (Chen et al.) 2018)). CNF mod- els have shown promising performance on generative tasks Kobyzev et al. (2020). However, known computational chal- lenge of CNF models is model regularization, primarily due to the non-uniqueness of the flow transport. To regularize the flow model and guarantee invertibility, [Behrmann et al. (2019) adopted spectral normalization of block weights that leads to additional computation. Meanwhile, (Liutkus et al. 2019) proposed the sliced- Wasserstein distance. Finlay et al. (2020); Onken al. (2021) utilized optimal tranppoortoortttto costs, and (Xu et al.) 2022) proposed Wasserstein-2 regularization. Although regularization is important to maintain invertibility for general-form flow models and improves performance in practice, merely
Figure 1: Comparison of IKO-iFlow (proposed) and other flow models. The JKO scheme approximates the transport of a diffusion process and the ResNet trained block-wise.

---

Under review as a conference paper at ICLR 2023
using regularization does not resolve non-uniqueness of the flow and there remains variation in the trained flow depending on initialization. Besides unresolved challenges in regularization, there remain several practical difficulties when training such models. In many settings, flows consist of stacked blocks, each of which can be arbitrarily complex. Training such deep models often places high demand on computational resources, numerical accuracy, and memory consumption. In addition, determining the flow depth (c.g., number of blocks) is also unclear.
In this work, we propose JKO-iFlow, normalizing flow network which unfolds the Wasserstein gradient flow via neural ODE invertible network, inspired by the JKO-scheme Jordan et al. (198)). 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. Each block in the flow model implements one step in the JKO-scheme can be trained given the previous blocks. As the JKO scheme pushes forwards the density to approximate the solution of Fokker-Planck equation of diffusion process with small step-size, the trained flow model induces smooth trajectory of density evolution, as shown in Figure [ The theoretical assumption does not incur a restriction in practice when training, whereby one can use larger step sizes coupled with numerical integration techniques. The proposed JKO-iFlow model can be viewed as trained to learn the unique transport map following the Fokker-Plann equation.
Unlike most CNF models where all the residual blocks are initialized together and trained end-to and-to-end, the proposed model allows a block-wise training which reduces memory and computational load. We further introduce time reparametrization with progressive refinement in computing the flow network, where each block corresponds to a point on the density evolution trajectory in the space of probability measures. Algorithmically, one can thus determine the number of blocks adaptively and refine the trajectory determined by existing blocks. Empirically, such procedures yield competitive performance as other CNF models with significantly less computation.
The JKO Flow approach proposed in this work also suggests a potential constructive approximation analysis of deep flow model. Method-wise, the proposed model differs from other recent JKO deep models. We refer to Section 1.1 for more details. In summary, the contribution includes
We propose neural ODE model where each residual block computes a JKO step and the training objective can be computed from integrating the ODE on data samples. The network has general form and invertibility can be satisfied due to the regularity of the optimal pushforward map that minimizes the objective in each JKO step.
We develop an block-wise procedure train the invertible JKO-iFlow network, which determines the number of blocks adaptively. We also propose a technique to reparametrize and refine an existing JKO-iFlow probability trajectory. Doing so removes unnecessary blocks and increases the overall accuracy.
Experiment wise, JKO-iFlow greatly reduces memory consumption and the amount of computa- tion, with competitive/better performance as several existing continuous and discrete flow models.
1.1
RELATED WORKS
For deep generative models, popular approaches include generative adversarial networks (GAN) (Goodfellow al. 2014 Gulrajani et al.) 2017: Isola et al.) 2017) and variational auto-encoder
Kingma & Welling 2014 2019). Apart from known training difficulties (e.g., mode collapse
(Salimans 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 al. 2019). Monge-Ampere flow (hhagg 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-Fuclidean data (Mathieu & 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 flow 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 & Ermon, 2019 Ho et al.] 2020)] Block et al. 2020 Song et al. 2021) under different setting. SSpecifically, these models estimate the
2

---

Under review as a conference paper at ICLR 2023
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 & Ermon (019). The recent work of Song et al.) (202T) developed reverse-time SDE sampling for score-based generative models, and adopted the connection neural ODE to compute the likelihood; using the same idea of backward SDF, /Zhagg] & Chen 202 1 proposed joint training of forward and backward neural SDEs. Theoretically, latent diffusion Tzen & Raginsky (2019b a) was used to analyze neural SDE models. The current work focuses on neural ODE model where the deterministic vector field f (2: t) is to be learned following a JKO scheme of the Fokker-Plann Planckk equation. Rather than neural SDE, our approach involves no sampling of SDE trajectories nor learning of the score function. Our obtained residual network is also invertible, which can not achieved by the diffusion models above. We experimentally obtain competitive or improved performance against on simulated and high-dimensional tabular data.
JKO-inspired deep models have been studied in several recent works. (Bunne et al. (022) reformu- lated 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. (021), which adopted input convex neural networks (ICNN) (Amos et al. ICNN as a special type of network architecture may have limited expressiveness (Rout et al. Korotin et al.) 2021). D. In addition to using gradient of ICNN, (Fan al. T02T) proposed to parametrize the transport in a JKO step by residual network but identified difficulty in calculating the push-forward distribution. The approach in (Fan et al.) 202T) also relies on variational formulation which requires training an additional network similar to the discriminator in GAN using inner-loops. In contrast, our method trains an invertible neural-ODE flow network which enables the flow from data density to normal and backward as well as the computation of transported density by integrating the divergence of the velocity field along ODE solutions. The objective JKO step to minimize KL divergence can also be computed directly without any inner-loop training, cf. Section 4
2017).
2022
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. (211)) established approximation by composition of Barron functions (Barron 1993); Bailey & Telgarsky (2018) developed space-filling approach, which was generalized in Perekrestenko et al. 2000 2021); Lu & 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 far, 4) that induces the Pokker-Planck equation, cff Section 3.2 Our model potentially leads to a constructive approximation analysis of neural ODE flow model to generate data density px-
2 PRELIMINARIES
Normalizing flow. A normalizing flow can be mathematically expressed via density evolution equation of p(x, t) such that p(z. 0) px and as increases p(x;, t) approaches pz ~ (0, Ta) Tabak & Vanden-Fijnden (2010). Given initial distribtuion p(z., 0), such flow typically is not unique We consider when the flow is induced by an ODE of æxt) in Idd
i(t) f(x(t),t). t),
(I)
where z0) ~ px. The marginal density of zt) denoted as p(x, t), and it evolves according to the continuity equation (Liouville equation) of [1 written as
d,p+ + V (pf) = O, p(x,0) = px
Ornstein-UUllnneck (OU) process. Consider a Langevin dynamic denoted by the SDE dX, = -VV(X,)dt + 2dW,, where V is the potential of the equilibrium density. We focus on the case of normal equilibrium, that is, V(x) = |x|? /2 and then pz VV In this case the process is known as the (multivariate) OU process. Suppose Xo ~ px. and let the density of X, be p(z, t) also denoted as p.(.). The Fokker-Planck equation describes the evolution of pt towards the equilibrium pz as
(3) Under generic conditions, Pt converges to pz exponentially fast. For Wasserstein-2 distance and the standard normal Pz, classical argument gives that (take C = 1 in Eqn (6) of Bolley et al. (2012)
a,p = V (PVV + Vp), V(x) == -1////2 p(z, 0) = (x).
W2(Pt. Pz) We2(P, P2), i 1>0 0.
3

---

Under review as a conference paper at ICLR 2023
JKO scheme. The seminal work Jordan et al. (1998) established time discretization scheme of the solution to (3) by the gradient flow to minimize (I(plipz) under the Wasserstein-2 metric in probability space. Denote by P the space of all probability densities on Rf with finite second moment. The JKO scheme at k-th step with step size h O, starting from p po E P, is written
1
p = arg min Flp] +
W (p
F(p) == KL lpz).
(5)
pEp
2/
was proved in Jordan all (1998) that h + O. p converges to the solution p( kh) of (3) for all k, and the convergence (=(),t) + p( is strongly in L (RR (0, T)) for finite T, where is piece-wise constant interpolated on (0, T") from p(t),
JKO SCHEME BY NEURAL ODE
Given i.i.d. observed data samples X, RR i ,,..., N, drawn from some unknown density px, the goal is to train an invertible neural network to transports the density px to an a priori specified density pz in Rd. where each data sample X, is mapped to code Z, A prototypical choice of pz is the standard multivariate Gaussian (0,, Ia). By slight abuse of notation, we denote by px and pz both the distributions and the density functions of data X and code respectively.
3.1 THE OBJECTIVE OF JKO STEP
We are to specify f(x; 1) in the ODE [ to be parametrized and learned by a neural ODE, such that the induced density evolution of p(x, t) converges pz & increases. We start by dividing
the time horizon [o, T into finite subintervals with step size h, let tx = kh and Ig+1 == 1 fx+1).
Define (x) p(x, kh), namely the density of z(t) at t kh. The solution of (I) determined by the vector-field f(x. 1) on L T++1 (assuming the ODE is well-posed (Sideris) 2013)) gives a one-to-one mapping Tk+1 Rd. s.t. Tk+ (a=())
x((4+1) and Tk+11 transports pę into Pk+11
ie.. (Pk) Pk. where we denote by Twp the push-forward of distribution p by T, such that (T+P)()) = (-)-
Suppose can find f(-,t) on I4+1 such that the corresponding Ta+1 solves the JKO scheme (5). then with small h. Pk approximates the solution to the Fokker-Planck equation 3 which then flows towards pz. By the Monge formulation of the Wasserstein-2 distance between p and q as W2 q) = minT:Tap=g Eawple T(x)12, solving for the transported density P& by (5) is equivalent solving for the transport Tw+1 by
T arg min F(T)
E 2h
Bowpe 1/2 T(x)
F[T] = KL (pz).
(6)
T-R--R"
The equivalence between (5) and (6) is proved Lemma A.I
gives that the value of [ can be computed from f(x, Furthermore, the following proposition
E Ix+1 only once pa determined by f(x, t) for 1 < 4.- The counterpart for convex function based parametrization of Ti was given in
Theorem of (Mokrov al.
2021), where the computation using the change-of- of-riial-- differs as we udopt an invertible neural ODE approach here. The proof is left to Appendix AA]
Proposition 3.1. Given pk. to a constant independent from f(x, on E Ik+1.
KL(T = Ex(ty) (vlaltatti)) - V f(x(s). + c.
(7)
By Proposition 3.1 the minimization (6 is equivalent to
min (flr.)))eeee Ex(tsrpp ((co(teti)) - Ax+1
1 f(r(s), s)ds + - [[((1)1) -1((t)))-=((1 -x(()))))(((())²)))))))) (8)
2h
where r(t++1) a(tt)) + pta plut f(a(s), Taking a neural ODE approach, we parametrize (f(x, he as a residual block with parameter 0k+1, and then 8 is reduced to minimizing over 0ct1- This leads to block-wise learning algorithm to be introduced in Section
4

---

Under review as a conference paper at ICLR 2023
3.2 INFINITESIMAL OPTIMAL f(2,,)
In each JKO step of (8) let Pk denote the current density, pz be the target equilibrium density. In this subsection, we show that the optimal in (8) with small l reveals the difference between score functions between target and current densities. Thus minimizing the objective (8) searches for neural network parametrization of the score function V Vlog Pt without denoising score matching as in diffusion- based models (Ho ct al. 2020; Song et al. 2021).
Consider general equilibrium distribution g with differentiable potential V. To analyze the optimal
pushforward mapping in the small h limit, we shift the time interval [kh,, (I + 1)h} to be [0, h] to
simplify notation. Then ( is reduced
(v(x(M)) - Lo V f(x(sss + 2/ [r((h) -x())))
1 2h
min (f(x.)))e
Em(0)~p
(9)
where s<h)- = =00++ f- f(x((s), s)ds. In the limit of - 0+, formally, r<h) x0)) (0) (((
O(h²). and suppose V of y is C2, V(x(h)) = (()(0)) + (>V()(()) [(x(0))
+ 0(12). For
any differentiable density P. the (Stein) score function is defined Sp V log p. and we have
vV = Sq. Taking the formal expansion of orders h. the objective in (9) is written as
Egpu (v(x) + (-s,(x))
(,) f(x,0) - V f(x,0) 0)
[f O)1 + 0(h2))
(10)
2
Note that Eg-p Vzx) is independent of f(x, t), and the (hh) order term in (10) is over f(rr, 0) only, thus the minimization of the leading term is equivalent to
min Expp f()-)(-0))
(T,
Tyf Sq Vf.
(11)
where T, is known as the Stein operator Stein 1972). The T,f in [1] echoes that the derivative of KL divergence with respect to transport map gives Stein operator (Liu & Wang 2016). The Wasserstein-2 regularization gives 2 regularization in [II.. Let L be the L space on (IR p(z (z)dz)), and for vector field v on Rd, v E (p) if [v " (z) < co. One can verify that, when both Sp and S, are in (p), the minimizer of (11 is
('( = s - Sp.
This shows that the infinitesimal optimal f(x, t) equals the difference of the score functions of the cquilibrium and the current density.
3.3 INVERTIBILITY OF FLOW MODEL AND EXPRESSIVENESS
At time the current density of ztt) is Pt. the analysis in Section 3.2 implies that the optimal vector field f(x, t) has the expression as
f(x: () Sq - = -VV V Vlgg Pt-
(12)
With this f(x, t), the Liouville equation (2) coincides with the Fokker-Planch equation (3). This is consistent with that JKO scheme with small i recovers the solution the Fokker-Planck equa- tion. Under proper regularity condition of V and the initial density Po. the r.h.s. of [2) is also regular over space and time. This leads to two consequences, in approximation and in learning: Approximation-wiin the regularity of f(x, t) allows to construct a k-th residual block in the flow network to approximate [f(x, f)} when there is sufficient model capacity, by classical universal approximation theory of shallow networks (Barron 1993] Yarotsky (017). The JKO-iFlow model proposed in this work suggests constructive proof of the expressiveness of the invertible neural ODE model to generate any sufficiently regular density px, which we further discuss in the last section.
For learning, when properly trained with sufficient data, the neural ODE vector field f(x, 1; 0k) will learn to approximate (12). This can be viewed as inferring the score function of Pt, and also leads to invertibilty the trained flow net in theory: Suppose the trained f(r, t; 02) is close enough to (2). will also has bounded Lipschitz constant. Then the residual block is invertible as long as the step size is sufficiently small, e.g. less than 1/L where L is the Lipschitz bound of f(x: t; 6&). In practice, 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 T
5

---

Under review as a conference paper at ICLR 2023
4 TRAINING OF JKO-IFLOW NET
4.1 BLOCK-WISS WISE TRAINING
Note that the training of (k: + 1)-th block in (8) can be conducted once the previous i blocks are trained. Specifically, with finite training data {X; = 2x(0)) the expectation Exit)~ps in (8) is replaced by the sample average over {i, (kh))- which can be computed from the previous k blocks. Note that for each given z(t) x(t)), both r(f2+1) and the integral of
. in (8) can be computed by numerical neural ODE integrator. Following previous works, we use the Hutchinson trace estimator (Hutchinson) 1989, Grathwohl et al.. 2019) to estimate
the quantity V f in high dimensions. Applying the numerical integrator in computing (8) we denote the resulting k-th residual block abstractly as So, with trainable parameters Ox. This leads to block wise training of the normal- izing flow network, as summarized in Algorithm found the generative performance JKO-iFlow may vary de- pending starting choices of l2, but simple choice such as & = k often yields reasonably good performance. We discuss further the initial selection of t in Appendix C.II Meanwhile, one can use any suitable termination criterion Ter(k) in line 2 of Algorithm|] [ In our experi- ments, we monitor the per-dimension W2 loss W2(T P&) as defined in (6). and terminate training more blocks if the per-dimension loss is below e. Lastly, the heuristic approach in line 5 of training a "free block" (i.e., block without the W2 loss) is to flow the push- purw-forward density PL closer to pz. where the former is obtained through the first L blocks and the latter denotes the Gaussian density at equilibrium.
Algorithm 1 Block wise JKO-iFlow training
1. Regarding input parameters,,
Require: Time stamps {fx J. training data, termi- nation criterion Ter and tolerance level c, max- imal number of blocks Lmax-
1: Initialize k = 1.
2: while Ter(k) > 6 and k S Imaxx
3: Optimize fer upon minimizing 8 with mini-batch sample approximation, given
[Le 1--1 Set h: + +.
4: end while
5: L f- k. Optimize fer.. using ( with =
00. (o Free block, no W22 regularization.)
Note that Algorithm ! Significantly reduces memory and computational complexity: only one block is trained when optimizing (8). regardless of flow depth. Therefore, one can use larger data batches and more refined numerical integrator without memory explosion. In addition, one can train each block for fixed number of epochs using either back proppaaaaoagnnnnn or the NeuralODE integrator (Grathwohl al. 2019 adjoint method). We found direct back peccpprppggttion enables faster training but may also lead to greater numerical errors and memory consumption. Despite greater inaccuracies, we observed similar empirical performances across both methods for JKO-iFlow, possibly due the block-wise training that accumulates fewer errors than a generic flow model composed of multiple blocks.
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 Pk. & = 1,... K residual blocks of the proposed normalizing flow network. The two techniques are illustrated in Figure Additional details can be found in Appendix B
,,,,,,,,, K, associated with the
Trajectory reparametrization. We empirically observe fast decay of the movements W (T Pkk pk); in other words, initial blocks transport the densities much further than the later ones. This is especially unwanted because in order to train the current block, the flow model needs to transport data through all previous blocks, yet the current block barely contributes to the density transport. Hence, instead of having & := kh with fixed increments per block, we reparametrize the values of t4 through an adaptive procedure, which is entirely based on the W2 distance at each block and the averaged W2 distance over all blocks.
Progressive refinement. To improve the probability trajectory obtained by the trained residual blocks, we propose a refinement technique that trains additional residual blocks based on time steps l obtained after reparametrization. In practice, refinement can be useful when the time increment 4--1 1 for certain blocks is too large. In those cases, there may exist numerical inaccuracies as the loss ( is computed over a longer time horizon. More precisely, we increase the number of JKO-steps parametrized by residual blocks, where in practice, we training C additional "intermediate" blocks for density transport between Py and Pk+1 at each k.

---

Under review as a conference paper at ICLR 2023
Probability Trajectory
W2 Movement
Ps
Figure 2: Diagram illustrating trajectory reparametrization and refinement. The top panel shows the original trajectory under three blocks via Algorithm[]] The bottom panel shows the trajectory under six blocks after reparametrization and refinement, which renders the W2 movements more even.
5 EXPERIMENT
We first generate based (wo-dimensional simulated samples. We then perform unconditional and conditional generation high-dimensional real tabular data. We also show JKO-iFlow's generative performance on MNIST. Additional details are Appendix
5.1 SETUP
Competing Methods and metrics. We compare JKO-iFlow with five other models, including four flow-based model and one diffusion model. The first two continuous flow models are FFJORD (Grathwohl al.) (219)) and OT-Flow (Onken et al.) (021). The next two discrete flow models are IResNet (Behrmann et al. 2019) and IGNN (Xu et al. 2022). which replaces the expensive spectral normalization in IResNet with Wasserstein-2 regularization to promote smoothness. The last diffusson model is the score- score-based generative modeling based on neural stochastic differential equation (Song al. 202T). which call it ScoreSDE for comparison. We primarily interested in two types of criteria. The first is the computational efficiency in terms of the number of iterations (e.g., batches that the model uses in training) and training time per iteration. Due the block-wise training scheme,
F 47!
X=F-12)
XY==FY
X-F
(a) True data
JKO-iFlow
(b) FFJORD 4.10c-4 8.57c-4
(c) OT-Flow 7.25e-4 1.32c-3
(d) W2 IResNet 3.34c-3 3.39e-3
(e) ScoreSDE 5.16c-4 1.09c-3
MMD[m]: 3.52e-4 T: 2.06e-4,
T: 2.96c-4, MMD[c]: 8.87e-4
2-
X-F-112
(f) Fractal tree
(g) Olympic rings
(h) Checkerboard
Figure 3: Two-dimensional simulated datasets. The generated samples X by JKO-iFlow in (a) are closer to the true data X than competitors in (b)-(e). Under the more carefully selected bandwidth via the sample-median technique, MMD[m] in (20) by JKO-iFlow also closer to the threshold 7 in (1 than others. (f)-(h) visualizes generation by JKO-iFlow on more examples.
Table I: Inversion error Exz-px || (T)(t)) x||2 of JKO-iFlow computed from sample average on test data, where To denotes the transport mapping over all the blocks of the trained flow network.
POWER
GAS
MINIBOONE
BSD300
Rose
Fractal tree
Olympic rings
Checkerboard
1,48e-5
1.58c-6
1.09c-6
1.53e-5
3.30c-6
3,58e-5
2.24c-6
3.07c-5

---

Under review as a conference paper at ICLR 2023
(a) True X 2.78e-3, MMD[m]: 2.29c-2
JKO-iFlow
(b) IGNN 2.60e-2
(c) True X
JKO-iFlow
(d) IGNN
T: 1.30e-2, MMD[m]: 3.65e-2
1.09e-1 1.53e-1
r: 3.90e-3, MMD[c]: 3.65e-2
3.11e-2
T: 1.69e-2, MMD[c]: 7.50e-2
Figure 4: Conditional graph node feature generation by JKO-iFlow and iGNN. We visualize the con- ditionally generated samples upon projecting down to the first two principal components determined by true X We visualize generation at two different values of Y.
Table 2:
Numerical metrics on high-dimensional real datasets. All competitors are trained after 10 times more iterations (i.e., batches), because their performance under the same number of iterations not comparable to JKO-iFlow. Complete results are shown in Table A.1
MMD
MMD
1.756-4
2346-4
T: 4,39p-4
7: 6.87e-4
JKO-iFoo
76K
8.21e-4
1,26e-3
JKO FFlow
112K
7.97c-4
1.0le-3
d
POWER
OT:Flow
FIOORD
W; IResNet
76X
904K
76K
5.69c-4
2.760-3
1.38-33
9.62e-4
1,98c
2.726-3
MINIBOONE 43
OT-FFlow
W2 IResNet
FHOORD
488K
112K
5.476-3
1.27e-2
1.0le-3
1.03e-3
1.04c-3
IRosNet
304K
2.49p-2
1.0k-3
IK0-iFlow
SconSDE
76K
76K
5.96c-
1,34e-3
1933-4
71
0.56c
2.83-l
1.79c-3
ScorSSIE
112K
1,35c-
4.200-3
9.63e-
1,03c-3
GAS
OT:Flw
6.09e-3
BSDS300
OT-Flow
JKO-HHlww
396K
396K
4.83-3
8.55e-2
8.14c-2
3.03e-3
Wg IResNot
304K
7.14c-3
1.50k-2
IResNet
990K
5.52e-1
6,88e-1
TRes Net
304K
3.260-3
2.736-2
IResNet
990K
5.42o-1
5.95e-
SeareSDE
76K
1,3163
1,456-
Score
396K
5.516-1
6.62c-1
we measure the number of iterations for JKO-iFlow as the sum of iterations over all blocks. Using this metric allows us to examine performance easily across models, under a fixed-budget framework in terms of batches available to the model. 'he second is the maximum mean discrepancy (MMD) comparison (Gretton et al. 2012, Onken et al.] (02T), which is way of measuring the difference between two distributions based samples. Additional details for MMD appear in Appendix C.4
We also report negative log-likelihood as an additional metric in Table A.I
Conditional Generation. Due to the increasing need for conditional generation, we also apply JKO-iFlow for conditional generation: generate samples based on the conditional distribution X|Y. Most existing conditional generative methods treat Y an additional input of the generator, leading to potential training difficulties. Instead, we follow the IGNN approach (Xu al. 2022). which also incurs minimal changes to training. Additional details are in Appendix C.5
5.2 RESULTS
3d compare JKO-iFlow with the competitors on non- conditional generation, where the subcaption indicates MMD values (20) under both bandwidths and the corresponding thresholds in (21). We omit showing IResNet with spectral normalization as it yields similar results W2 IResNet. The generative quality by JKO-iFlow is the closest to that of the ground truth, and when the MMD bandwidth is more carefully selected via the sample-median technique, JKO-iFlow also yields smaller MMD than others. Meanwhile, Figures 3h shows the satisfactory generative performance by JKO-iFlow compares the performance of JKO-iFlow before and after using the technique described Section 4.2 where the generative quality by JKO-iFlow improves after several reparametrization moving iterations, and Figure A.4 shows additional unconditional and conditional generation results.
Two-dimensional toy data. Figures 3a
other examples. In Appendix, Figure A.2
High-dimensional tabular data. In terms of conditional graph node feature generation. Figure compares JKO-iFlow with IGNN on the solar dataset introduced in iGNN. The results show that JKO-iFlow yields competitive or clearly better MMD values on the conditional distribution X|y with the most or second most observations, respectively. Next, Table assesses the performance of JKO-iFlow and competitors on four high-dimensional real datasets, where JKO-iFlow still has the best overall performance. Dataset details are described in Appendix CC.3 To ensure fair comparison, we keep the number of parameters for continuous flow models and the diffusion model the same and omitted
properly increase model sizes for discrete flow models. FFJORD results on BSDS300 due to incomparably longer training time. In terms of results, except on POWER, JKO-iFlow yields smaller or very similar MMD under both bandwidths than all other methods. Table A.I in Appendix
8

---

Under review as a conference paper at ICLR 2023
Projected
Projected (2)
000
sudon
21
Wort:: IIS." AtR00 AdAlls
(b) Results moving iteration
24
Projected X
Projected F (-)
70
10
Moving Iteration
Bee
Iter
Ber
Bee
Iee
(a) Components of loss over moving iterations.
(c) Results moving iteration 5.
Figure 5: MINIBOONE, reparametrization moving iterations of JKO-iFlow. We plot different components of the loss objective (8) over tg. In (a), results moving iteration are obtained by using A (modified for training flow model) 4 times, and the reparametrization gives more uniform W2 losses after moving iterations. On this example, the generative performance are both good before and after the moving iterations, ef. plots (b) and (c).
19
799
9
4
O
5
9
7
7
9
0
9
7
8
5
2
0
6
4
5
3
5
6
3
8
3
3
0
7
3
4
2
3
5
1
1
8
6
2
4
8
2
6
8
6
3 9
Figure 6: MNIST generation by JKO-iFlow coupled with a pre-trained auto-encoder.
Cishoss the complete results, including the number of training iterations and test log- oigllikelhooo.. Overall, we remark that comparisons using MMD[m] (i.e., MMD with bandwidth selected using the sample-median technique) best align with visual comparisons in Figure A. of Appendix C.6 that we suggest MMD[m) as a more reliable metric out others we used. Furthermore, we illustrate the reparametrization technique on MINIBOONE in Figure 5 where the benefit appears in yielding flow trajectory with more uniform movement under a competitive generative performance.
MNIST. We illustrate the generative quality of JKO-iFlow using AutoEncoder. Consider a pre- trained encoder Enc: 787 + Rd and decoder Dec : R2 - 784 such that Dec (En(()) z X for flattened image X. We choose d = 16. The encoder (resp. decoder) uses one fully-connected layer followed by the ReLU (resp. Sigmoid) activation. Then, JKO-iFlow is trained on N encoded images (Eno(X,)) )= and the trained model gives an invertible transport mapping (over all residual blocks) To : Rdd + j The images are generated upon sampling noises Z (O, La) through the backward flow followed by the decoder, namely Dec(T, (Z)). The generated images are shown in Figure
DISCUSSION
The work can be extended in several directions. The application to larger-scale image dataset 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 toy data shown in this work. For the methodology, the time-continuity the parametrization of the residual blocks (as a result of the smoothness of the Fokker-Planck flow) have not been exploited in this work, which may further improve model capacity as well as learning efficiency. Theoretically, the model expressiveness of flow model to generate any regular data distribution can be analyzed based on Section 3.3 To sketch a road-map, a block block-wise approximation guarantee of f(z, t) in (12) can lead to approximation of the Fokker-Planck flow (3), which pushes
forward the density to be e-close to normality in T = log(1/e) time, cf. (4). Reversing the time of the ODE then leads to an approximation of the initial density Po px by flowing backward in time from to zero. Further analysis under technical assumptions is left to future work.

---

Under review as a conference paper at ICLR 2023
REFERENCES
David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. Optimizing functionals on the space of probabilities with input convex neural networks. arXiv preprint arXiv:2106.00777 2021.
Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pp. 146-155. PMLR, 2017.
Bolton Bailey and Matus Telgarsky. Size-noise tradeoff' in generative networks. Advances in Neural Information Processing Systems, 31, 2018.
Andrew R Barron. Universal approximation bounds for superpositions of sigmoidal function. IEEE Transactions on Information theory, 39(3):33-9945 1993.
Jens Behrmann, Will Grathwohl, Ricky TQ Chen, David Duvenaud, and Jörn-Heriik Jacobsen. Invertible residual networks. In International Conference on Machine Learning, pp. 573-582. PMLR, 2019.
Adam Block, Youssef Mroueh, and Alexander Rakhlin. Generative modeling with denoising auto- encoders and langevin sampling. arXiv preprint arXiv:2002.001107, 2020.
François Bolley, Ivan Gentil, and Arnaud Guillin. Convergence to equilibrium in wasserstein distance for fokker-planck equations. Journal of Functional Analysis, 263(6))2430-2457, 2012.
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal
transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 6511-6528. PMLR. 2022.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
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.
Laurent Dinh, Jascha Sohl- Dickksien,,,,, and Samy Bengio. Density estimation using real nvp. ArXiv, Abbs/6605008800 2017.
Jiaojiao Fan, Amirhossein Taghvaci, and Yongxin Chen. Variational wasserstein gradient flow, arXiv preprint ar arXiv:21 2021.
Matthias Fey and Jan F. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learming on Graphs and Manifolds, 2019.
Chris Finlay, Jöm-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How train your neural ode: the world of jacobian and kinetic regularization. In International conference machine learning, pp. 3154-3164. PMLR, 2020.
lan Goodfellow, Jean Pouget-Abadia Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial nets. In NIPS, 2014.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Kristjanson Duve- naud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. ArXiv, abs/1810.01367, 2019.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alex Smola. A kernel two-sample test. Mach. Learn. Res., 13:723-773, 2012.
Ishaan Gulrajani, Faruk Ahmed, Martín Arjovsky, Vincent Dumoulin, and Aaron C. Courville. Improved training of wasserstein gans. In NIPS, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-66555 68511 2020.
Chin- Men-Wei Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective diffusion- based generative models and score matching. Advances in Neural Information Processing Systems, 34: 2021.
10

---

Under review as a conference paper at ICLR 2023
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- 1989.
Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A. Efros. Image-to-image translation with con- ditional adversarial networks. 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 5967-5976, 2017.
Rie Johnson and Tong Zhang. A framework composite functional gradient methods for generative adversarial models. IEEE transactions pattern analysis and machine intelligence, 43(1):17-32, 2019.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker-planck equation. STAM journal on mathematical analysis, 29(1):1-17, 1998.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. CoRR, abs/1312.6114, 2014.
Diederik P. Kingma and Max Welling. An introduction to variational autoencoders. Foundations and
Trends@ in Machine Learming, 12(4):307-392, 2019. doi: 10.156//2200000056
Ivan Kobyzev, Simon JD Prince, and Marcus Brubaker. Normalizing flows: An introduction and review of current methods. IEEE transactions on paltern analysis and machine intelligence, 43 (11):3964-3979, 2020.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference Learning Representations, 2021. httptttss///ppeeeeeeeeeee wttp://eeeooo.edeooo.o.oo
Shampine Lawrence. Some practical runge-kutta formulas. Mathematics of Computation, 46:
135-150, 1986.
Holden Lee, Rong Ge, Tengyu Ma, Andrej Risteski, and Sanjeev Arora. On the ability neural nets to express distributions. In Conference on Learning Theory, pp. 1271-1296. PMLR, 2017.
Jenny Liu, Aviral Kumar, Jimmy Ba, Jamie Kiros, and Kevin Swersky. Graph normalizing flows. Advances in Neural Information Processing Systems, 32, 2019.
Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances neural information processing systems, 29, 2016.
Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian Robbr- Stöter. Slised-wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. In International Conference on Machine Learning, pp. 4104-4113, PMLR, 2019.
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.
James Lucas, G. Tucker, Roger B. Grosse, and Mohammad Norouzi. Understanding posterior collapse in generative latent variable models. DGS@ ICLR, 2019.
Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 33:2503-2515, 2020.
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.
Derek Onken, 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.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. Advances in neural information processing systems, 30, 2017.
11

---

Under review as a conference paper at ICLR 2023
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 Vito,, 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' d'Alché-Bal E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019.
Dmytro Perekrestenko, Stephan Müller, and Helmut Böleskei. Constructive universal high- dimensional distribution generation through deep relu networks. In International Conference on Machine Learning, pp. 7610-7619. PMLR, 2020.
Dmytro Perekrestenko, Léandre Eberhard, and Helmut Böleskei. High-dimensional distribution generation through deep neural networks. Partial Differential Equations and Applications, 2(5): 44, 2021.
Fredrik Ronquist, Maxim Teslenko, Paul van der Mark, Daniel L. Ayres, Aaron E. Darling, Sebastian Höhna, Bret R. Larget, Liang Liu, Mare A. Suchard, and John P. Huelsenbeck. Mrbayes 3.2: Efficient bayesian phylogenetic inference and model choice across a large model space. Systematic Biology, 61:539 542, 2012.
Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference Learning Representations, 2022. URL htttps://pperrevieet
Tim Salimans, Ian J. Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. ArXiv, tass/6066034999 2016.
Thomas C. Sideris. Ordinary differential equations and dynamical systems. 2013.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations, 2021.
Charles Stein. A bound for the error in the normal approximation to the distribution of sum of dependent random variables. In Proceedings of the sixth Berkeley symposium on mathematical statistics and probability, volume 2: Probability theory, volume 6, pp. 583-603. University of California Press, 1972.
Esteban Tabak and Eric Vanden-Eijnden. Density estimation by dual ascent of the lieliioll Communications in Mathematical Sciences, 8(1):217-233) 2010.
Belinda Tzen and Maxim Raginsky. Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit. arXiv preprint 1:Xivv1905.09: 2019a.
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory, pp. 3084-31 14. PMLR, 2019b.
Chen Xu, Xiuyuan Cheng, and Yao Xie. Invertible neural networks for graph prediction. arXiv preprint MXXX::226600163. 2022.
Dmitry Yarotsky. Error bounds for approximations with deep relu networks. Neural Networks, 94: 103-114, 2017.
Linfeng Zhang, Lei Wang, et al. Monge-ampls cre flow for generative modeling. arXiv preprint arXiv:1809.10188. 2018.
Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in Neural Information Processing Systems, 34:16280-16291, 2021.
12

---

Under review as a conference paper at ICLR 2023
A PROOFS
A. PROOFS IN SECTION
Lemma A.l. Suppose p and q are two densities on I in P, the following two problems
min Lolp] (L(plla) pEp
W (p,p).
(13)
2h.
1
min TB
[r([T] -TT.(TTTll0 +
E Ezupl | -T(x)1?
(14)
2h
have the same minimum, and (a) 1fT TRd + Rd is a minimizer of (14). then p² = (T* # is a minimizer of
13
(b) If p" is minimizer of (13). then the optimal transport from p to minimizes (4
Proof of Lemma A.1
A.1
Let the minimum of 14 be Li and that of (13 be Lip-
Proof of (a): Suppose LT achieves minimum 7", then T the optimal transport from p to
P² (T #p because otherwise Lr can be further improved. By definition of Lp. we have Lplp"] 2 L. We claim that L L. Otherwise, there another p' such that L,lp" < T.j. Let T be the optimal transport from p to p', and then LT [T'] = [ppp/] L.p.
Lij = Lr
contradicting with that 1. is the minimum of Lr. This also shows that Tipl0"] Top = 1. that is, p" a minimizer of I.p-
Proof of (b): Suppose Lp achieves minimum at p². Let T be the OT from p to p", then Ez-pl T* (x)) = W2(P, p²)² and then Lr[T"] = Lp(p"] LÃ which equals L7 as proved in (a). This shows that 7 a minimizer of LT.
Given pa being the density of zx()) at f = kh, recall that T' is the solution map from (t) to r +++ h). We denote P: == Pk. and PLth == TTp:. By definition,
Proof of Proposition 3.1
KL
Ex-peth (log Plog Peth (z) log Pzx))....
(15)
Because pz x e V (x) /2/ we have logpz(x) Vxx) + ci for some constant ei. Thus
Ex-penn log Pz (x) = (t (z h)) = E (a((t h)).
(16)
To compute the first term in 15 note that
E Exoeper log (z) = Ex(()=o, log Pith (a h)).
(17)
and by the expression (called "instantaneous change-of-varine variable formula" normalizing flow litera- ture (Chen et al. 2018), which we derive directly in below)
d
log p(x(t). V f(((t))).
(18)
dt
we have that for each value of a(t).
1+h V. f(x(s))
log Pet (x(( + h)) = log p(x(t + h), t + h) = log (())),)) -)
Inserting back to 17 we have
H+h lthh V ())ss)
Exepiia log ( E2(())<p log pe(x((() Ex(rr
The first term is determined by Pe = Pk. and thus is a constant independent from f(x, t) on t [kh, (k 1)4). Together with (16), we have shown that
++h
r.h.s. of 15) C2 Ex(e)
f(x(s), s)ds) + Ex(e)~p, V(a(t + h)),
which proves 7
13

---

Under review as a conference paper at ICLR 2023
Derivation of (88 by chain rule,
d
log logg(s(t), vp -(t) + (sp(a(t).t) ((x(t),t)
dt
Vp. v--f-v-(0 (pf)) P (a(0).))
(by and
- V f(((tt),
B TECHNICAL DETAILS OF SECTION 4.2
Although the layer-wise training formulation in Section enjoys several aforementioned benefits, there exists undesirable movement patterns along the trajectory. Empirically, the movement by initial blocks fo, is much larger than later ones. The blue curve labeled "Phase 1" in Figure A.2a visualizes one typical pattern of the movement measured by W22 distances.
In fact, this phenomenon arises due to smaller gradient magnitude at later estimates, which gradually approach a local minimum during optimization. particular, such irregular movement also appears in gradient descent vector space. We thus propose aegarametrize-and-retee technique.
specific training flow networks by the JKO scheme. It essentially
Vector-space case
We first motivate method with optimization in vector space. Suppose goal is to find local minimum x² of Fzz) for nonlinear differentiable function F IR + R. Starting at (0)), consider the following sequential optimization problem, where (!()) denotes the estimate at the I-th iteration and h pre-specified regularization parameter:
(t+1)) arg min F(x) + 2h || x.')
(19)
Using the first order Taylor expansion F(x) =e + VF(x()) xl1) at x(1), we get
+
(1+1) = x()) higt, 9: = v=F(1(1))
Define the are length of iterates S, ||,.((+11 x011) ||2 = 1el1lge11. whereby it appears in practice that the magnitude of S, is near zero (1) + x². This issue typical result of small gradient as estimates approach the local minimum. We thus propose Algorithm 2 to resolve this uneven are length issue, which takes in iterates (ffoold) and step sizes hold from the previous trajectory.
We first motivate and explain the reparametrizztion step in line 3. Mathematically, lengths defined using re-optimized values (f.neww) to satisfy Spew ~ S. This is equivalent to requiring hy "eg""" ||2 & S, where gitw : (=(peew). The quantity [[g2"w [| is unknown before re-optimization takes place, so that we approximate using |ا = S (t.ol). In practice, using the quantity Sheld/ /Sold alone to update h; can be undesirable. because larger heew tend to cause non-smooth trajectories and inaccurate final estimates. We thus introduce inertia controlled by parameter 77 and upper bound the largest h by hmax to allow more flexibility.
want
We now explain the refinement step in line 4. We interpolate 2 intermediate points between each pair of (z((t.ew), (((w+.,eww) For instance, if C 1, we optimize for the "mid-point" ((++1/2,uew before reaching (t+ww, Using this approach ensures smoother new trajectories (ar(t.eww)) )221 and potentially more accurate final estimate. Figure (.33 illustrates the behavior and our solution on minimizing the Muller-Ban Brown energy potential in 2
JKO Flownet reparametrization
Although Algorithm 2 is developed for re-parametrizing and refining trajectories in vector space Rd it can be directly used to reparametrize for JKO-iFlow by replacing the are length S, between consecutive iterates in vector space with the W2 movement in probability space of the residual
14

---

Under review as a conference paper at ICLR 2023
Algorithm 2 Trajectory improvement (vector-space case)
Require: Penalty factors helh and iterates x (t.old)) for I..... T. and 17 E (0, 1].
hmax
Compute Soldd : ||x: [1=(+++ x((,old) 2| and S = : Spldd /T.
2: for t= .,.., T do
3: Compute hg = min (1ighd + n7 (Shh / nold), Mmax)- {o Reparametrize)
4: For C 2 1, store hnow == [hgew /(C + 1)..................... hh.....,ee, /(C + 1)] for C' + times. {D Refine]
5: end for
6: Re-optimize 19 for (f..new) with (hrew Uoo.
7: Repeat all steps above until std([Sp })/ mean mean((Sp ((pma)) small enough.
block. More precisely, let L be the total number of trained blocks via Algorithm [ and denote hy = 4+11 - l/ as the "step-size" for block fe.- Replace the iterates x(() in vector space with a(tx), which is the mapping through previous W2 distance, which can be easily computed using N samples trajectory. The refinement step thus becomes training additional residual blocks via optimizing (8).
blocks. Then, the are length S, becomes the
along each step of the
C EXPERIMENTAL DETAILS
C.1 CHOICE OF 4 IN ALGORITHM
Recall that to train our JKO-iFlow, one needs as input sequence of tæ, where the k-th JKO block integrates from io fx+1- Although the selection of i varies by problem, we consider two choices in our settings.
Constant increment. Denote l4 := |x+1 lg. We let hg = ci for a constant Ci > 0. On many experiments for vro-dimensional toy data and high-dimensional data,
C1 1. Constant multiplier. Given tu > and a constant C22 we let 4x+1 := c2t2. The rationale is that from empirical evidence, the W2 movement as in (6) tends be larger at initial blocks than at latter blocks, so that moving later blocks more than the initial ones would enable more uniform movements, thus faciliating the training process. On some experiments for two-dimensional toy data and high-dimensional data, we let to 0.75 and c2 = 1.2.
We acknowledge that many other choices are possible. We also want to emphasize that due to the reparametrization and refinement techniques proposed in Section 4.2 the values of te would be adaptively updated based on data, where the adaptive values would yield more uniform W2 movements over blocks as we saw in Section
C.2 OTHER SETUP DETAILS
All experiments are conducted using PyTorch (Paszke et al. 2019) and PyTorch Geometric (Fey &
Lenssen 2019) Regarding network architecture
For simulated 2D data, nigh-dimensional real data, and MNIST using pre-trained autoen- coder: each residual block uses fully-connected layers of the form + + H + where d (resp. Hf) is the feature (resp. hidden nodes') dimension. The hidden dimension vary by example, in the range of 128-5522.
For conditional graph node feature generation: each residual block uses one Chebnet input layer of order 3 followed by two fully-connected layers. The hidden dimension H = 64 in all hidden layers.
The activation function is chosen as Tanh Softplus with 20. In addition, we use Rugge-Kutte 4 (Lawrence 1986) to numerically estimate the integrals in the continuous flow model. The Adam optimizer is used with constant learning rate of 1e throughout training. Regarding batch sizes, we use 1000 samples for simulated two-dimensional data and 50% of the training samples for the solar graph data. The batch size selections for nigh-dimensional real data described in Table A.1
15

---

Under review as a conference paper at ICLR 2023
C.3 DATASET
For two-dimensional simulated examples, we generate fresh random draws of 10000 training samples each training epoch. The four high-dimensional real datasets (POWER, GAS, HEP- MASS, MINIBOONE) come from the University of California Irvine (UCI) machine learning data repository. These datasets are commonly used compare flow models (Grathwohl et 2019 Finlay al. 2020 Onken et al., 202T). The solar dataset as used in iGNN (Xu et al. 2022) is retrieved from the National Solar Radiation Database (NSRDB).
C.4
MMD METRICS
Besides visual comparison, the maximum mean discrepancy (MMD) Gretton et al. 2012: Onken (et al. 2021) provides a quantitative way to evaluate the performance generative models. Given j==1 and a kernel function k(2, v), we compute
samples X = { and Y == {yz {yi}
N
M
M
N
MMD(X, x, Y)
Σ
k(Vi U5)
k(a, (20)
N
M
NM
1=1
w11² with bandwidth h. We select the bandwidth both as constant value he = 2 and via the "sample-median technique" (Gretton et al. 2012) hm 2median( [I|x; 2y 22 where x, test samples. We use the same set of 200 test samples to compute hm in each experiment. We thus denote MMD[c] (resp. MMD(m]) as evaluating 20) using the Gaussian kernel with constant (resp. median) bandwidth, where X (resp. Y) denotes true (resp. generated) test sample. Note that a low MMD value indicates two samples X and Y are likely drawn from the same distribution (Gretton al. 2012). In this setting, MMD an impartial evaluation metric as it is not used to train JKO-iFlow or any competing methods.
For our purpose, use the Gaussian kernel k<<x,, y) =-
We also determine the statistical significance of MMD value. First, compute the threshold Q 1- (1MMD( [ x (
(21)
where Qi-a denotes the upper - a quantile of set of scalars and Z cl,...,, .N} denotes the j-th index at the b-th bootstrapping without replacement. Then, under the null hypothesis that X and Y are drawn from the same distribution, this hypothesis is rejected if MMD exceeds the threshold т. The Type-I error is controlled at level a. Thus, if the MMD values by two models both exceed r, we prefer the model with the smaller MMD. If both values are under т, then they generate equally well. our experiments, we use 1000 bootstraps, each of which has 50% re-sampled test samples.
C.5 CONDITIONAL GENERATION
We follow the conditional generation scheme as proposed in iGNN (Xu et al. 2022). More precisely, when the response variable Y is categorical variable taking value in K classes, iGNN designs the target distribution as a Gaussian mixture model. Thus, instead flowing from data density px to noise density pz. iGNN flows from the conditional data density Pxly PHIY, where H Y ~ (HY, o²1). One can then minimize the negative log-likelihood loggxx>> using logPHIY and the change-of- variable formula.
To use JKO-iFlow for conditional generation in this setting, we thus only need to modify the objective Instead of using Vz based on Z N(O, Ia), we would using VHLY based the Gaussian mixture H H [[pyyo21]
C.6 ADDITIONAL RESULTS
We present complete results in addition to those in Section 5 In particular,
Table A.I 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., (222]). which computes the evidence lower bound (ELBO) for the data log likeelildoood reported in the last column. In addition, Table A.2 A.22 contains MMD and negative log- likeliiod results for FT-Flow and FFJORD as taken from the original papers.
16

---

Under review as a conference paper at ICLR 2023
Table A.I: Numerical metrics on high-dimensional real datasets, in addition those Table Comparing to flow-based models, JKO-iFlow takes much less iterations to reach a small enough MMD value. Although ScoreSDE the fastest, its performance, even under 100 times more iteration than JKO-iFlow, is still worse in terms of MMD[m] on all except GAS. We advocate the comparison using MMD[m] because the results align with visual comparisons in Figure A.I
Model
Training
Testing
Time (h)
lter
Time/lter (S)
Batch size
MMD ml
MMD[c]
Neg Loglik
1.75e-
2.84c-4
JKO-sFlow
76K
0.12
0.76K
10000
8.21644
1.26c-3
0.58
POWER
OT-Flow
FFJORD
76K
76K
1.21
7.58K
7.58K
0.57
100000
5.09c-4
9.62e-4
0.30
W2 IResNet
304K
0.49
7.58K
0.23
100000
2.76--3
2.720-3
0.36
IRosNet
304K
0.76
7.58K
0.36
10000
4.50c-3
2.49e-2
3.37
ScorSDEE
76K
7.58K
1.346-3
6.56e-3
3,41
ScoreSDE
76K
0.84
73.85K
0.04
10000
1,30e-3
5.66e-3
3,33
JKOwFFlww
76K
0.08
0.76K
0.38
5000
5.96e4
1.79e-3
-4.61
OT-Flow
76K
0.72
7,60K
0.34
5000
1.51e-3
3.64e-3
-4.9
d
W2 IResNet
FFJORD
76K
304K
0.57
3.49
7.60K
7.60K
0.27
1.65
5000
5000
3.62e.3
7.14e.3
6.09e-3
1,50e22
-4.5
-2.07
3.26e-3
ScoreSDE
76K.
0.04
7.60K
0.02
5000
1,31e-3
1.45e-3
-3.69
SD
76K
0.42
76.00K
0.02
5000
4.27e-4
8.56e
-5.58
4.59c4
7:
6.87c-4
KOw-Flow
112K
0.03
0.32K
0.33
2000
7.97e.4
1.01e33
13.63
MINIBOONE d' 13
112K
3.39K
1.01e-3
11.93
RFORD
112K
1.74
3.39K
1.85
2000
5.47e33
1.04e-3
23.45
|ResNet
3.25K
1.276-2
1.0cc3
16.34
IResNet
448K
1,32
3,25K
1.46
2000
2.58e-3
L.4e-3
22.36
SeoreSDB
112K
0.01
3,25K
0.01
2000
4.29>-3
1.03e-3
27.38
SooreSDR
1.10e-3
20.70
1,35e4
#: 9.63e-5
-15.667
BSDS300 63
OT-Flow
396K
3.50
10.29K
1.22
1000
8,553-2
8.44c-2
-142.45
W3 IResNet
99UK
2.01
10,29K
0.70
1000
5,52e-1
6.88e-1
-10.339
IReaNet
990K
3.47
10.29K
1.21
1000
5.42-1
5.95e-1
-33.11
SooreSDE
396K
0.01
10.29K
0.005
1000
5,516-1
6.62e-1
-7.55
6.65e-1
Table A.2: MMD[c] and negative loglikelihood results of OT-Flow and FFJORD, as taken from
(Onkn et al. 2021). We include them to compare against ours in Table The models in previous studies use comparable model size (especially for OT-Flow), where the numerical results in some cases are much smaller than ours due significantly longer training time.
Training
Testing
Data Set
Model
Param
Time (h)
Iter
MMDJe]
Neg Loglik
POWER
OT-Flow
18K
22K
4.68e-5
-0.30
d 6
FFJORD
43K
68.9
29K
4.34e-5
-0.37
GAS d=
OT-Flow
127K
6.1
52K
2.47e-4
-9.20
FFJORD
279K
75.4
49K
1.02e-4
10.69
MINIBOONE
OT-Flow
78K
2.84c-4
10.55
d == 43
FFJORD
821K
9.0
16K
2.84c-4
10.57
BSDS300 d = 63
OT-Flow
297K
7.1
37K
4.24e
-154.20
FFJORD
6.7M
166.1
18K
6.52e-3
-133.96
Figure A.
visualizes the principal component projections of the generated samples by JKO-iF low and competitors of the high-dimensional real datasets.
Figure A.2
visualizes components of loss 8 and the resulting generated images.
visualizes the trajectory of estimates in I2a of minimizing the Muller-Broo energy potential.
Figure A.3
Figure A.4 shows additional unconditional and conditional generated samples by JKO-iFlow on toy data.
17

---

Under review as a conference paper at ICLR 2023
Projected
Projected (
Projected F-12))
F-)
(2)
Projected 22
(a) Power: JKO-iFlow
(h) OT-Flow
(c) FFJORD
(d) Wa IResNet
(e) ScoreSDE
Projected
Projected (-)
Projected F-)
Projected (-1
Projected F-1
Projected F 1z)
(f) Gas: JKO-iFlow
(g) OT-Flow
(h) FFJORD Projected Az!
W2 IResNet Projected F 72))
i Score
Projected x
Projected (2)
Projected 4Z)
Projected z())
(k) MINIBOONE: JKO-iFlow
OT-Flow Projected F-)
(m) FFJORD
lRes
ScoreSDE
Projected
Projected PrajcctedF-1iz F-)
Projected F FlZz)
Projected F-
Projected (1/2
(p) BSDS300: JKO-iFlow
(q) OT-Flow
(r) W2 IResNet
(s) IResNet
(L) ScoreSDE
Figure A.I: Generative quality 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
Divergence: -f. slds
X
X=F-1
[& 1. ((0G))) stels11B/22
-11," 6.x.st)
(b) Results at moving iteration
6.6
X
-=
Moving
iteration
Iter
Iter
Ieer
Iter
(a) Losses moving iterations.
(c) Results at moving iteration 4.
Figure A.2: Rose, reparametrization moving iterations of JKO-iFlow. The plots and setup are identical Figure 5 We observe improved generative quality after the moving iterations.
18

---

Under review as a conference paper at ICLR 2023
16 arc lengths at (moving iteration: mean & std)
reparam-1:0.344 & 0.097
teparam-9:0.357 & 0.018
refine-4:0. 1:0.184 & 0.009
reparam-2:0.34 2:0.34-2000.77 0.078
reparam-10:0.380 10:00337 & 0.015
refine-5:0.188 5:0.184 0.007
9eparam-3:0.355 0.062
1:paaam-11:0.357 & 0.013
refine-6: 6:0.184 0.006
reparam-4:0.352 0.051
reparam-12:0.355 0.010
refine-7:0.184 0.004
reparam-5:0. 0.041
reparam-133 3.8 0.008
refine-8:0 saaee-8:.1184 0.003
reparam
refine 9:0.184 0.003
reparam-7:0..35
refine-2:0 2:0.183
refine-10:0.184 0.002
reparam-8:0..d 0.022
refine-3:0.184 0.012
refine-11:0.184 0.002
0.5
0.4
0.2
0.20
0.25
0.30
0:35
0.45
time 2h:
After proximal steps, max inner loop steps Estimated minimum is [-1.911 0.105]
3.0
2.5
2.0
1.5
1.0
0.5
8.0
16: (-1.94 0.244]
0.5
The minumum: (-1.911, 0.103)
-100
5.0
4.5
4.0
-3.5
-30
2.5
2.0
15
Figure A.3: Reparametrization and refinement moving iterations vector space based on Algorithm The task is to estimate local minimizer of the Muller Brown energy potential. We see that are lengths between consecutive iterates become more even in magnitude over more reparametrizztiion and refinement moving iterations.
2-F 4z)
X -
All -(XI)
(a) Two-cireles
(b) Conditional generation illustration.
Figure A.4: Additional unconditional and conditional generation simulated toy datasets by JKO-iFlow.
19


Reviews:
[{"review_id": "wwVnyykeKm", "reviewer": "Reviewer_T8Fn", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper under consideration proposes a normalizing flow model (called JKO-Flow) which combines JKO scheme and Continuous Normalizing Flows. The previous approaches which dealt with JKO scheme parameterized the pushforward transform $T(x)$  arising in JKO objective either by ICNNs [1, 2] or general NNs of the form $F : \\mathbb{R}^D \\rightarrow \\mathbb{R}^D$ [2] . The main idea of the authors is to replace the pushforward $T(x)$ with continuous dynamics $T(x) = x + \\int\\limits_{t_k}^{t_{k + 1}} f_k(x(t)) d t$ modeled via the techniques developed for Continuous Normalizing Flows.\n\n[1] Mokrov et. al. Large-Scale Wasserstein gradient flows, https://arxiv.org/pdf/2106.00736.pdf\n\n[2] Fan et. al., Variational Wasserstein gradient flows, https://arxiv.org/pdf/2112.02424.pdf\n\n**Post-rebuttal:** I increase my score by 1", "summary_of_the_review": "The authors present a new idea which combines CNFs and JKO scheme. However, there are some questions regarding the usefulness of the proposed approach. Moreover, It is recommended to stress-test the method on the image generation problems which are standard benchmarks for many modern Normalizing flows models.\n", "strength_and_weaknesses": "**Strength**\n\n- To the best of my knowledge, the proposed idea is fresh. It seems that the continuous dynamic parameterization helps to leverage some difficulties faced by previous approaches (like $\\log \\det J_{T}$ estimation [1] or $\\min \\max$ objective [2]). Given a good integral estimator the learned JKO-Flow could be easily inverted in a neural ODE manner without need of specific optimization procedures [1]. \n- The idea to optimize block-composed Neural Network by sort of \"divide and conquer\" block-wise algorithm also seems to be interesting.\n\n**Weaknesses**\n- No comparison with  [2]. The proposed JKO-Flow approach seems to be directly related to their method. \n- It would be interesting to see the generative properties of the JKO-Flow approach directly on data space of nontrivial image datasets like CIFAR10 / CelebA. As I understand, the proposed approachcan be straightforwardly adopted for such (image-data) scenarios.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "qBcZQlBIkDX", "reviewer": "Reviewer_tb5k", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes using Wasserstein gradient flow of KL dviergence  to construct the flow model. ", "summary_of_the_review": "See above", "strength_and_weaknesses": "First of all, the paper has a bad definition of notations. For example, the definition of V (is the potential of the equilibrium density) appears in the preliminaries. I think the paper should organize a problem-setting section to make all things clear.\n\nHowever, in objective (8) there is a V (is the potential of the equilibrium density) included. How to get the V for general distribution? V includes fitting a score, which is very similar to yang's ICLR paper. \n\nThe main concern is the realtionship to paper [1,2]. [1] consider the gradient flow of f-divergence and gives out a flow with velocity field f''*nabla (p/p_t). Thrid, how this paper differs from [3]. The only difference is the objective for training but not the gradient flow aiming to approximate. ( If the author claim their objective is better, I need to see evidence)\n[1] https://arxiv.org/abs/2012.06094\n[3] Johnson R, Zhang T. A framework of composite functional gradient methods for generative adversarial models[J]. IEEE transactions on pattern analysis and machine intelligence, 2019, 43(1): 17-32.\n\nIn terms of the literature review, this paper only discusses how they are different in terms of motivation (or how to derive the model), but have no discussion of how the derived model are different from each other. \n\nMissing experiments on CelebA. Missing baseline using diffusion model.\n\nMissing reference:\n[1] Liutkus A, Simsekli U, Majewski S, et al. Sliced-Wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions[C]//International Conference on Machine Learning. PMLR, 2019: 4104-4113.\n[2] Zhang L, Wang L. Monge-amp\\ere flow for generative modeling[J]. arXiv preprint arXiv:1809.10188, 2018.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Q8n6PPMqYS", "reviewer": "Reviewer_ffZu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This current paper proposes a normalizing flow algorithm that implements the JKO scheme using neural ODE flow blocks. At each time interval, an optimization over a vector field parameterized as ResNet is done to obtain the velocity field that can be integrated to obtain samples at the next step. The normalizing flow is invertible and the objective functional of the JKO scheme is taken to be the KL divergence with respect to a standard Gaussian so that by inverting the flow we can generate the data distribution. The proposed method appears to be more efficient than the alternatives as demonstrated by the experiments.\n", "summary_of_the_review": "Overall I think the paper lacks novelty and needs more comparison with other methods (other JKO methods, diffusion models) to demonstrate its effectiveness. The writing can also be improved.\n", "strength_and_weaknesses": "## Strengths:\n* Optimizing the velocity field of the ODE for JKO instead of transport maps that are more common in the existing literature is novel.\n* The idea of using one block for each JKO step is interesting and it seems to have resulted in efficiency.\n\n\n## Weaknesses:\n* I found the novelty of the proposed method somewhat limited. The only difference from [Alvarez-Meliset et al. 2021], [Mokrov et al. 2021] is that instead of parameterizing the pushforward map as ICNN (note in either work it's not necessary to use ICNN; arbitrary networks can also be used, following the same reasoning at Lemma A.1), the current work uses neural ODE. I'm not convinced by the superiority of using neural ODE over a pushforward map, which the current work does not compare against. In my opinion, a more severe problem is not addressed: namely at step $k$ all these algorithms need to push initial samples by $k$ steps to obtain samples for the current iteration, which scales overall quadratically in $k$.\n* The section 2.2 is not new to my knowledge and references are missing. See a similar derivation in Theorem 3.1 of Liu et al. \"Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm\".\n* I appreciate that two heuristic enhancements are proposed in Section 3.2. However, at first glance, these two enhancements seem contradictory. On the one hand, trajectory reparameterization wants to remove blocks, but on the other hand, progressive refinement requires adding intermediate blocks for large time steps. It would be great to have some ablation studies illustrating the effectiveness of either enhancement. \n* Since the target application is generative modeling, and the overall idea follows diffusion models (namely mapping the data distribution to a standard Gaussian, and then inverting the process), I think the authors should compare with diffusion models as well. \n* The writing is a bit sloppy overall. There is almost no assumption stated for Lemma A.1 or Proposition 2.1. At the very least we probably need $p$ and $q$ to have a finite second moment along with differentiability assumptions. The proof of 2.1 does not seem novel to me. It is essentially a direct application of the instantaneous change of variable from [Chen et al. 2018] (which is by itself just rewriting the continuity equation). There is missing reference on this formula (see the equation under (17)), and moreover as written this formula is not correct. The derivative needs to be a total derivative, i.e., it should be $d/dt (\\log \\rho(x(t), t)) = -\\nabla \\cdot f(x(t), t)$. The sloppiness of the writing is also manifested in the many handwavy sentences in the main text. To give a few examples:\n    - Above (4), \"Under generic conditions\" --- what conditions?\n    - In the second paragraph of 2.1, \"The solution of (1) ... gives a one-to-one mapping\" this should only be true if the time interval is small enough, by Picard–Lindelöf theorem\n    - In the last paragraph of page 5, \"... it will also have bounded Lipschitz constant\". How is this true? There is no assumption on $\\rho_t$ having bounded the Lipschitz constant.\n    - At the end of page 5, \"The analysis is postponed here\", where more details could have been given\n* For the conditional generation experiments, more details will be helpful. I don't understand what it means by \"we evaluate $V$ for a Gaussian mixture $H | Y$\".\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "I9CRgBcGFcD", "reviewer": "Reviewer_y2DU", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a continuous-time normalizing flow model based on the JKO scheme for gradient flows in the Wasserstein-2 space. A neural ODE flow network is trained to approximate the JKO scheme. This results in a block-wise training procedure where each block approximates a JKO step. Several such blocks are trained sequentially and the final flow is obtained via a composition of the blocks. Experiments demonstrate that the proposed method performs competitively against baseline models while being more computationally efficient. ", "summary_of_the_review": "The paper presents a technically interesting idea of training continuous normalizing flows using the JKO scheme. The proposed method results in improved computation efficiency during training. However, this work has issues with the presentation, justification of claims, and quality of empirical evaluation. ", "strength_and_weaknesses": "[Strengths]\n* The paper presents an interesting method of training a continuous-time normalizing flow to approximate the JKO scheme. This allows the model blocks to be trained sequentially with potentially larger batch sizes. \n* While the JKO scheme has been explored before in the machine learning context for large-scale gradient flows (Mokrov et al., 2021), the proposed framework is novel to the best of my knowledge.\n \n[Weaknesses]\n* The quality and clarity of the write-up and the overall presentation can be improved.\n    * The need for the discussion in Section 2.2 is unclear. The result in Eq. (12) can be directly arrived at by using the fact that JKO scheme recovers the FPK equation which gives us the steepest descent curve. Then a direct matching of the terms in FPK for the OU process with the Liouville Equation gives  the result in Eq. (12). Also see the discussion in Appendix D of Song et al., 2021.\n    * Section 3.2 can be better explained. I did not understand what \"progressive refinement\" is and how exactly does it help the overall training.\n    * The scores can be reported in a more readable format. The \"e\" notation is messy.\n    * The paper ends abruptly after the MNIST experiment without a follow-up discussion or conclusion.\n    * Minor issues:\n        * What is meant by \"trajectory obtained by **rest blocks**?\"\n        * Rephrase \"can be **trained by iteslf**\".\n        * \"Th theoretical assumption\" -> \"The...\"\n        * \"space of probability\" -> \"space of probability measures\".\n        * The \"T\" for termination criterion in Algorithm 1 can be replaced with a different symbol. Currently it coincides with the \"T\" for the optimal transport map.\n        * A few other minor typos here and there. Please proofread.\n* On the claims in introduction:\n    * The authors discussed non-uniqueness of the flow as a motivation. However, it is unclear how/if the flow induced by the proposed method is unique. How does minimization via Eq. (8) ensure uniqueness? Just empirical evidence is not enough. Previous methods have also provided such evidence.\n    * It's unclear if the motivation of pruning unused blocks is actually relevant to other models. This seems to be a issue with the proposed model wherein a naive implementation results in unused blocks.\n* On experiments:\n    * The experiment results are not convincing enough. Most of the setups considered in the paper are toyish in nature. \n    * Why do the MMD results in Table 2 differ significantly from the results reported in Onken et al., 2021? If this is due to a difference in setups, what happens if the setup described in Onken et al., 2021 is used for experiments? What motivated the choice of specific bandwidth in MMD? In my experience, MMD is very sensitive to bandwidth choices. Can a better metric be used?\n    * Is it possible to evaluate the log-likelihood under this model? If yes, why did the authors not compare the log-likelihoods of the models?  \n* On the relation to prior work:\n    * With the specific choice of the OU process and the discussion around Eq. (12), the relation to probability flow ODE in Song et al., 2021 should be better discussed. \n\n**Post Rebuttal Edit**: Increased score from 3 to 5.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject