Normalizing flow neural networks by JKO scheme
Authors: Chen Xu
Xiuyuan Cheng
Yao Xie
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Normalizing flow
invertible neural networks
JKO scheme
TL;DR: Neural ODE model which learns the deterministic transport equation that solves the Fokker-Planck equation of the diffusion process, allowing block-wise training inspired by the JKO scheme at reduced computation and memory cost.

Abstract: Normalizing flow is a class of deep generative models for efficient sampling and likelihood estimation, which achieves attractive performance, particularly in high dimensions. The flow is often implemented using a sequence of invertible residual blocks. Existing works adopt special network architectures and regularization of flow trajectories. In this paper, we develop a neural ODE flow network called JKO-iFlow, inspired by the Jordan-Kinderleherer-Otto (JKO) scheme, which unfolds the discrete-time dynamic of the Wasserstein gradient flow. The proposed method stacks residual blocks one after another, allowing efficient block-wise training of the residual blocks, avoiding sampling SDE trajectories and score matching or variational learning, thus reducing the memory load and difficulty in end-to-end training. We also develop adaptive time reparameterization of the flow network with a progressive refinement of the induced trajectory in probability space to improve the model accuracy further. Experiments with synthetic and real data show that the proposed JKO-iFlow network achieves competitive performance compared with existing flow and diffusion models at a significantly reduced computational and memory cost.

Normalizing flow neural networks by JKO scheme
Chen Xu School of Industrial and Systems Engineering Georgia Tech
Xiuyuan Cheng Department of Mathematics Duke University
Yao Xie School of Industrial and Systems Engineering Georgia Tech
Abstract
Normalizing flow is class of deep generative models for efficient sampling and likelihood estimation, which achieves attractive performance, particularly in high dimensions. The flow often implemented using sequence of invertible residual blocks. Existing works adopt special network architectures and regularization of flow trajectories. this paper, we develop a neural ODE flow network called JKO- iFlow, inspired by the Jordan-Kinddn -eee--ninnneeeer--OOOO the discrete-time dynamic the Wasserstein gradient flow. The proposed method stacks residual blocks after another, allowing efficient block-wise training of the residual blocks, avoiding sampling SDE trajectories and score matching or variational learning, thus reducing the memory load and difficulty end-to-end training. We also develop adaptive time reparameterization of the flow network with progressive refinement of the induced trajectory in probability space improve the model accuracy further, Experiments with synthetic and real data show that the proposed JKO-iFlow network achieves competitive performance compared with existing flow and diffusion models at significantly reduced computational and memory cost.
(JKO) scheme, which unfolds
Introduction
1221--100
Generative models have wide applications in statistics and machine learning to infer data- generatierrnggg distributions and to sample from the model distributions learned from the data. In addition to widely used deep generative mod- els 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 great potential. The flow model learns the data distribution via an invertible mapping F between the data density px in Rd and the target standard multi- variate Gaussian density Pz. Z ~ (0 I,). While flow models, once trained, can be utilized for efficient data sam- pling and explicit likelihood evaluation, training of such models often difficult in practice. To alleviate such diffi- culties, many prior works DDinh et al.
(a) JKO :Flow
(b) usual CNF
Figure I: Comparison of JKO-iFlow (propesed) and stan- dard CNF models. contrast most existing CNF els, JKO-iFlow leams the unique deterministic transport equation corresponding to the diffusion process by directly performing black wise training of neural ODE model.
2015 2017 Kingma 2019 Ruthotto al.
and Dhariwal 2018 Behrmann al.
2020 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
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

models have shown promising performance on generative tasks Grathwohl et al. (2019), Kobyzev et al. 2020]. However, 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 such as FFJORD Grathwohl et al. (2019) may follow less regular trajectory in the probability space, see Figure lb), 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. (020). Xu et al. 2022] Huang et al. [2023], only using regularization may not resolve ve the non-uniqueness of the flow. Besides regularization, several practical diffieulties 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 high demand on computational resources and memory consumption.
In this work, we propose JKO-iFlow, an invertible normalizing flow network that unfolds the Wasserstein gradient flow via neural ODE model inspired by the JKO-scheme [Jordan et al. [998]_
The JKO scheme, see (5), can be viewed as proximal step to minimize the Kullback-Leible (KL) divergence between the current density and the equilibrium density. It recovers the solution of the Fokker-Planck equation (FPE) 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 FPE which flows from the data distribution toward the normal equilibrium and gives smooth trajectory of density evolution, see Figure Ta). Unlike most CNF models, where all the residual blocks are trained end-to-end, each block in the JKO-iFlow network implements one step in the JKO scheme to learn the deterministic transport map by minimizing an objective of that block given the trained previous blocks. The block-wise training significantly reduces memory and computational load. Theoretically, with small step size, the discrete-time transport map approximates the continuous solution of FPE, which leads to the invertibility of each trained residual block. This leaves the residual blocks to be completely general, such as graph neural network layers and convolutional neural network layers, depending on the structure of data considered in the problem. The theoretical need for small step sizes does not incur a restriction in practice, whereby one can use step size not exceeding a certain maximum value when adopting the neural ODE integration. We further introduce time reparameterizatioo 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. The algorithm adaptively chooses the number of blocks and step sizes.
The proposed approach is related to the diffusion models (Song and Ermon 2019 Ho et al.. 2020 Block et al., 22200 (ong et al., 2021] yet differs fundamentally in that our approach is a type of flow models, which directly computes the data likelihood, and such likelihood estimation is essential for statistical inference. While the diffusion models can also indirectly obtain likelihood estimation, they are more designed as samplers. In terms of implementation, our approach trains a neural ODE model without SDE sampling (injection of noise) nor learning of score matching. It also differs from previous works on progressive training of ResNet generative models JJohnson and Zhang (2019). Fan et al. [2022) in that the model trains an invertible flow mapping and avoids inner loops of variational learning. We refer to Section 1.1 for more discussions 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 pz having a parametrized potential V. exemplified by the application to the conditional generation setting where pz is replaced with Gaussian mixtures Xu et al. 2022]. In summary, the contributions of the work include
We propose an invertible neural ODE model where each residual block corresponds to a JKO step, and the training objective can be computed from pushed data samples through the previous blocks. The residual block has a general form, and the invertibility is ensured due to the regularity and continuity of the approximate solution of the FPE.
We develop block-wise procedure to train the JKO-iFlow model, which can adaptively determine the number of blocks. We also propose to adaptively reparameterize the computed trajectory in the probability space with refinement, which improves the model accuracy and the overall computational efficiency.
2

---

We show that JKO-iFlow greatly reduces memory and computational cost when achieving com- petitive or better generative performance and likelihood estimation compared to existing flow and diffusion models on simulated and real data.
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)| 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. 20200 has been extensively developed, including continuous flow (Grathwohl et al.) 2019]. Monge-Amenn Ampere flow
/Zhang et al. 2018). discrete flow [Chen et al. 2019]. and extension to non-Euclidean data [Liu et al.
2019 Mathieu and Nickel, [020 Xu et al. 2022]. Efforts have been made to develop novel invertible
mapping structures [Dinh et al. 2017 P et al. [017]] and regularize the flow trajectories
by transport cost (Finlay al.]2222 Onken et al., [22 I Ruthotto et al.. 2020 Xu et al.] 2022
Huang et al., 2023). 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., (021] under different settings. Specifically, these models estimate the score function (i.e., the 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 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 compute the likelihood; using the same idea of backward SDE, Zhang and Chen [021] proposed joint training of forward and backward neural SDEs. Theoretically, latent diffusion Tzen and Raginsky [2019baa was used to analyze neural SDE models. The current work focuses on a neural ODE model where the deterministic vector field f(z, t) be learned from data following a JKO scheme of the FPE. 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, example.[Song et al. [2021) derived neural ODE model from the learned score function of the diffused data marginal distributions for all 1. We experimentally obtain competitive or improved performance against the diffusion model on simulated two-dimensional and high-dimensional tabular data.
IKO-inspired deep models have been studied in several recent works. Bunne et al. 2022 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
al. 2022 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 jet al., 2022 Korotin et al. 2021]. In addition to using the gradient of ICNN, (Fan et al., 2022 proposed parametrizing the transport in a JKO step by residual network but identified difficulty in calculating the push-forward distribution. The approach variational formulation, which requires training an additional network simil lrr to the discriminator in GAN using inner loops. The idea of progressive additive learning in training generative ResNet, namely training ResNet block- by a variational loss, dates back to Johnson and Zhang [019] 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 direetly without any inner-loop training, see Section 4
[Fan et al. 2022) also relies on a
Compared to score-based neural SDE methods, our approach is closer to the more recent flow-based models related to diffusion models [Lipman et al. 2023 A Albergo and Vanden-Eijnden 2023 Boffi and Vanden Eijnden| 2023). These works proposed to learn the transport equation (a deterministic ODE) corresponding to the SDE process. Specifically, Lipman et al. 2023 Albergo and -ijndenEijnden
3

---

2023] matched the velocity field from interpolated distributions between initial and terminal ones; Bofh and Vanden-Eijnden 2023) proposed to learn the score log Pe (z) (where P. solves the FPE and is unknown a priori) to solve high-dimensional FPE. Using Stein's 's identity (which is equivalent to the derivation in Section (.2). the step-wise training objective (Bofh and Vanden-Eijnden, [023]
optimizes to learn the score without the need to simulate the entire SDE. The idea of approximating
the solution of FPE by a deterministic transport equation dates back to the 90s Degond and Mustieles
[1990]. , Degond and Mas-Gallic [1989), and has been used in kernel-based solver of FPE in [Maoutsa ct al. 2020] and studied via a seff-consistency equation in [Shen et al. 2022]. While our approach learns equivalent velocity field at the infinitesimal time step, see section 3.2 the formulation in JKO-iFlow is at a finite time-step motivated by the JKO scheme. We also mention that an independent concurrent work (Vidal et al. [023) proposed similar block-wise training algorithm using the JKO scheme under the framework of [Onken et al. 2021 and demonstrated benefits in avoiding tuning the penalty hyperparameter associated with the KI-ddivegeence objective. Our model is applied to generative tasks of high-dimensional data, including image data, and we also develop additional techniques for computing the flow probability trajectory.
For the expressiveness of deep generative models, approximation properties of deep neural networks for representing probability distributions have been developed in several works. Lee et al. [017]]
established approximation by composition of Barron functions [Barron [993]: Bailey and Tel garsky
[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 metries, 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. Our JKO-iFlow model potentially leads to a constructive approximation analysis of the neural ODE flow model.
2 Preliminaries
Normalizing flow. A normalizing flow can be mathematically expressed via a density evolution equation of p(z,t) such that p(z, 0) = px and as increases p(z, t) approaches pz A((0, La) (Tabak and Sanden-Eijnden 2010]. Given an initial distribution p(z, 0), such a flow typically is not unique. We consider when the flow is induced by an ODE of xxt) in IRd
=tt) = (x)(t),
(I)
where z0)) ~ px. The marginal density of 2 (t) is denoted p(ar, t), and it evolves according to the continuity equation (Liouville equation) of [1 written as
Oep + V (pf) = O, p(z, 0) = (x))..
(2)
CynsteinUUhlenbed (OU) process. Consider a Langevin dynamic denoted by the SDE dX,, = -V,(X,) + V2dW,, where V is the potential of the equilibrium density pz- We focus on the case of normal equilibrium, that is, V(x) = |x: 2 /2 and then pz e In this case, the process is known as the (multivariate) OU process. Suppose Xo ~ px. and let the density of X, be p(x, f) also denoted as pc(.). The Fokker-Planck equation (FPE) describes the evolution of P: towards the equilibrium pz as follows, where V(x)= == /2/
ap= = V (DVV + Vp), p(z, 0) (z).
(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)
W2(Pe.PZ) & e "W2(PoPZ2)
1>0 0.
JKO scheme. The seminal work Jordan et al. 1998 established a time discretization scheme of the solution to 33 by the gradient flow minimize KL [Pz) under the Wasserstein-2 metric in probability space. Denote by P the space of all probability densities on Rd with finite second moment. The JKO scheme computes a sequence of distributions Pk. k = 0, starting from Po = Po E P. With step size h > 0, the scheme at the k-th step is written as
1
Pk+1 arg min Flpl + PEP
W (pkk p).
(5)
2h
where F'c] := KL(pllpz). It was proved in Jordan al. [1998] that as + 0, the solution pk converges to the solution p(-, kh) of (3) for all k, and the convergence (1m)(-,)) - p(s t) is strongly in L (Rd, (O, for finite T where P(n) is piece-wise constant interpolated from pk.

---

JKO scheme by neural ODE
Given i.id. observed data samples X; R L,..,, NN drawn from some unknown density px- the goal is to train an invertible neural network to transport the density px to an priori specified density pz in Rd, where each data sample X; is mapped to code Z, A prototypical choice of pz the standard multivariate Gaussian (0, Ia). In this work, we leave the potential of pz abstract and denote by V. that is, pz x e-V and V(=) = |x|²/ /2 for normal pz. By slight abuse of notation, we denote by px and pz both the distributions and the density functions of data X and code Z respectively.
,.....,,,
3.1 Objective of JKO step
We are to specify f(x f) in the ODE (I) to be parametrized and learned by a neural ODE, such that the induced density evolution of p(x, t) converges to pz as increases. We start by dividing the time horizon [0, T into finite subintervals with step size h, let tx- = kh and Ix+1 == [ fx+1).
Define (z) : p(x:, kh), namely the density of z(t)
kh. The solution of (T) determined
by the vector-field [(x () on I E Tx+1 (assuming the ODE is well-posed [Sideris 2013) )) gives
one-to-one mapping Tx+1 on Rd s.t. Tx+1(x(tx)) = x(fx+1) and Th+1 transports pk into Pk+1+
i... (Tk) (P=)) = Pk. where we denote by Tw the push-forward of distribution p by T. such that (A) pTT (A)) for measurable set A. In other words, the mapping T is the solution
map of the ODE from time ly to 1++1-
Suppose we can find f(:()) on T4+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
as (p q) = minn:TTp== 1 T(x)1 solving for the transported density ps by (5) is equivalent to solving for the transport Tx+1 by
T Ta+1 = arg min
F[T) + Ex-pu ||:e T(x)
(6)
T:Rd SRd
2h
where F[2]- KL Ipz). The equivalence between and 66 proved in Lemma A.1
Furthermore, the following proposition gives that, once Pk determined by f(x, t) for S tg- the value of F[T] can be computed from f(x, t) on t E I++1 only. The counterpart for convex junction-based parametrization of Th was given in Theorem of [Mokrov et al., 2021]. where the computation using the caaune-of-vvrrallie differs as adopt invertible neural ODE approach here. The proof is left to Appendix A
Proposition 3.1. Given pks up to a constart e independent from f(x, 1) on E la+1.
<L<T lPz) = En(t))<p (v(o(tet)))
2x+1 V f(x(s), s))s) + c.
(7)
By Proposition 3.1 the minimization is equivalent to
1 1 [a(t)+1) - xtxx))))
min Ex(t))-p (V (x((+++)) (f(x:0))eett
V f(I(s))
+
(8)
2h
where r(f2+1) =(ta) + fli plat ((z(s), Taking neural ODE approach, we parametrize (f(x, ()) as residual block with parameter 0 and then ( is reduced to minimizing over 8x41. This leads to block-wise learning algorithm to be introduced in Section where we further allow the step-size h to vary for different & as well.
3.2 Infinitesimal optimal f(x, 1)
In each JKO step of 8 let Ps denote the current density, Pz be the target equilibrium
density. this subsection, we show that the optimal (8) with small I 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 log pr implicitly, in contrast
to diffusion-baai models which learn the score function explicitly (Ho et al. 2020
Song al.
5

---

2021], e.g., via denoising score matching. At infinitesimal h, this is equivalent to solving the FPE by learning deterministic transport equation in [Both and Vanden-Eijnden 2023 Shen et al. 2022]
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, (A+ 1)/} to be [0, h] to simplify the notation. Then ( is reduced to
(V(x(h))) - ((r()))), ((-o)))))) s)ds + 2/ [><(()) - x00]]]²]
1
min (f(2.)))
(9)
2h
where (hh) x00) f(r(()), s)ds. In the limit of h 0+, =00 (f))))))))((()))))))))))))))))))
O(h²), and suppose V of q is C2, ((x(h)) ((x(0)) + NVV ((x((0)))) + ( For
any differentiable density the (Stein) score function is defined as Sp V log P. and we have
VV Sq. Taking the formal expansion of orders of h, the objective in (9) is written as
EEEO ((()) + (-s,(x)) f(x,0) V f(xr,)) + [f(x, 0)1² +0(+)
(10)
Note that E Egp V(x) is independent of fiz, (). and the O(h) order term in (T)) is over f(r 0) only, thus the minimization of the leading term is equivalent to
min Bgwp (()=((-0)
To
Tyr:= Sy ff
(11)
where Tg is known as the Stein operator [Stein 1972]. The 7, in 11 echoes that the derivative of KI. divergence with respect to transport map gives Stein operator Liu and Wang 2016). The Wasserstein-2 regularization gives an La regularization in (11. Let (p) be the L space on
(R", p(z)dr),, and for vector field v on V E (2(p) if [v<)) < co. One can verify that, when both Sp and Sy are in (p), the minimizer of 11 is
(-. Sq Sp.
This shows that the infinitesimal optimal f(x, equals the difference between the score functions of the equilibrium and the current density.
3.3 Invertibility of flow model and expressiveness
At time the current density of (t)) Pt, the analysis in Section 3.2 implies that the optimal vector field f(I,t) has the expression as
f(z, t) = S. Sp. -V V log Pt.
(12)
With this f(x, t), the Liouville equation JKO scheme with a small h recovering the solution to the FPE. Under proper regularity condition of V and the initial density Pa: the r.h.s. of (12) is also regular over space and time. This leads to two consequences, in approximation and in learning: Approximationn-wise, the regularity of f(x; I) allows construct k-th residual block in the flow network to approximate (f(rr, (e) when there is sufficient model capacity, by classical universal approximation theory of shallow networks [Barron] 1993 Yarotsky [2017]. We further discuss the approximation analysis based on the proposed model in the last section.
2 coincides with the FPE 33. This is consistent with the
For learning, when properly trained with sufficient data, the neural ODE vector field f(x, t; 0k) will learn to approximate (12). This can be viewed as inferring the score function of Pt, and also leads to the invertibility of the trained flow net in theory: Suppose the trained f(, t: 0k.) is close enough to [1): it will also have bounded Lipschitz constant. Then the residual block invertible as long as the step size h is sufficiently small, e.g. less than l/L where L is the Lipschitz bound of f(x, I; 0s). In practice, we typically use smaller i than needed merely by invertibility (allowed by the model budget) so that the flow network can more closely track the FPE of the diffusion process. The invertibility of the proposed model is numerically verified in experiments (see Table A.I
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) illustrated in a vector space in Appendix B.3

---

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; z, 2,(0)]== the expectation Ex(t)=p= in (8) is replaced by the sample average over (k)]) which can be computed from the previous
blocks. Note that for each given æ(() = x=(4), both a((x++1) and the integral of V 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 compute the quantity V f in high dimensions, and we also propose a iiite-difference approach to reduce the computational cost (details in Appendix B.2). Applying the numerical integrator in computing (8), we denote the resulting k-th residual block abstractly as fo. with trainable parameters 8/..
This leads to a block-wise training of the normalizing flow network, as summarized in Algorithm T The sequence of time stamps la is given by specifying the time steps h2s == (a+1 lg. which we allow to differ across k. The choice of the sequence hg is initialized by geometric sequence starting from ho with maximum stepsize hmmax: see Appendix B.1 is one, the sequence of h4 gives a constant step size. The adaptive choice of h1p with refinement (by Regarding the termination criterion Ter in line 2 of Algorithm[] [ we monitor the ratio Ex-p, I - (z) /Lm-p 11 |2 and terminate when it is below some threshold e, set as 0.01 in all experiments. In practice, when training the k-th block, both the averages of || (x))1 and (x)]l² are computed by empirically averaging over the training samples (in the last epoch) at no additional computational cost. Lastly, line of training a "free block" (i.e., the block without the W22 regularization) is to flow the gush-forward density PL closer to the target density pz. where the former obtained through the first L. blocks.
In the special case where the multiplying factor
adding more blocks) will be introduced in Section 4.2
The block-wise training significantly reduces the memory and computational load since only one block is trained when optimizing (8) regard- less 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 empiri- cally observe that training each block using stan- dard pookppropagttion (rather than the adjoint method in neural ODE) gives comparable re- sult at lower cost. To ensure the invertibility of the trained JKO-iFlow network, we further break the time interval [t&-1, (x) into or 5 subinter- vals to compute the neural ODE integration, e.g., by Runge-Kutta-4, Kutta-4, We empirically verify small inversion errors on test samples.
Algorithm 1 Block-wiss JKO-iFlow training
Require:
Time stamps {& J. training data, termi- nation criterion Ter and tolerance level e, max- imal number of blocks Lmax-
1: Initialize k: 1.
2: while Ter(k) > c and ls < Lmax do
3: Optimize for upon minimizing ( with mini-batch sample approximation. given
1-1 [ Set ks t +
4: end while
5: L + k. Optional: Optimize for.. without
W2 regularization.
4.2 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,..., associated with the residual blocks of the proposed normalizing flow network. The two techniques are illustrated in Figure 2 Further details and illustrations of the approach can be found in Appendix B
Trujectory reparameterization We empirically observe fast decay of the movements (T Pk, P&) when hi 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-Plann flow, see ( but unwanted in the algorithm because in order to train the current block, the flow model needs to transport data through all previous blocks, and yet the later blocks trained using constant step size barely contribute to the density transport. Hence, instead of having constant fgs we reparameterize the values of tj through adaptive procedure based on the W₂ distance at each block. The procedure uses an adaptive approach to encourage the We movement in each block to be more even across the L blocks, where the retraining of the trajectory can be potentially warm warm-started by the previous trajectory in iteration.
7

---

Progressive refinement. The performance of CNF models typically improved with larger number of residual blocks, corresponding to smaller step size. The smallness of stepsize also ensures the invertibility of the flow model, in theory and also in practice. However, directly training the model with small non-adaptive stepsize may result in long computation time and convergence to the normal density qz only after a large number of blocks, where the choice of h2 is not efficient after adaptive reparam- eterization. We introduce a refinement approach to increase the number of blocks progressively, where each time interval ( 4x) splits into two halves, and the number of blocks doubles after the adaptive reparameterization converges 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 h2, which benefits the accuracy of the JKO-iFlow model including the numerical accuracy in integrating each neural ODE block.
Probablity Irajectory
W; Movement
Figure Diagram illustrating trajoctory reparamcterization and refine-
ment. Top: the original trajectury under three blocks via Algoritmm Bottom: trajectory under refinement, which renders the W2 movements more even.
reparameterization
the trajectory
5 Experiment
In this section, examine the proposed JKO-iFlow model simulated and real datasets, including both unconditional and conditional generation tasks. Codes are available at hhtttss htop:::///tthu/ com/ hamrel- cxu/ JK0-iFlow
5.1 Baselines and metrics
We compare 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- 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). The diffusion model baseline is the score-based neural SDE [Song et al. [021], which we call "ScoreSDE." Details about the experimental setup, including dataset construction and neural network architecture and training can be found in Appendix C.2
The accuracy of trained generative models is evaluated by two quantitative metrics, the negative ligeiikll (NLL) metric, and the kernel maximum mean discrepancy (MMD) [Grettnn al.. 2012a metric, including MMD-1, MMD-m, and MMD-c for constant, median distance, and custom bandwidth respectively. The test threshold t is computed by bootstrap, where an MMD metric less than indicates that the generated distribution is evaluated by the MMD test to be the same as the true data distribution (achieving a 0.05 test level). See details in Appendix [C.T] The computational
cost is measured by the number of mini-batch stochastic gradient descent steps and the training
JKO-iPlow 2.73c-4
(b) FFJORD 3,88c
(c) OT-How 1.42e-3
ScorSSDE
r: 2.79c-4, MMD-c
3.14c-3
6.90c-4
NLL
2.64
2.95
3.30
3.35
3.2
(f) Fractal tree 3.12e-4, MMD-c:
(g) Olympic rings
(b) Checkerboard
2.17c-4 2,20
7: 3.16e-4, MMD-e:
2.36c-4 1.66
3.99-4, MMD-c:
2.70c-4
NLL
NLL
NLL
3.59
Figure 3: Results two dimensional simulated datasets by JKO iFlow and competiturs.
8

---

Table I: Resalts tabular datasets. All competioors trained fixed hudget selup using 10 times mini batches (their performances using the same number of mini-batches are worse and not comparable to JKO-iPlow). See the complete table in Table A.2
The results using
more computation are given in Table A
Dais Sat Model
Param
Test bMM
Text MMD-1
NLT
Data Set
Model
Piram
Test m
Tea MMD-1
NLL
FI L.73-4
7: 2-30c-4
12 2.446--4
1: 3.75c-4
IKO-illow
76K
9.86c-5
2.40p-4
-0.12
IKO-
112K
9,66e-4
3,7%
12.55
POWER
OT:
IGNN
HEDRD
76K
314K
7,56c-1
1.93-43
1.16e-3
5.35p-4
1.59/3
0.63
0.32
0:95
MINIBOONE
d 13
IGNN
PRODRD
488K
112K
112K
3.5
6. 6.58e-4
1.21e
4.01e
3.70-4
4.12e4
23.77
26.45
11.44
314K
392:-3
2.46-22
428K
2.13c-3
4.160-
22.56
ScoreSDI
9.12r-
6.8e-3
SconSDDI
4.35e-
27.38
1355-4
273cc4
1,36-4
T:
1.86-4
JKO- ilow
76K
1.52c 4
5.00e 4
7.65
JKO FFlow
396K
2.24e-4
1.9c
-13.82
GAS
HIDRD
Flow
76K
76K
1887-3
1.9
3.28e-3
5.16e-4
-2.65
6.04
BSDS300 d' 63
CTT-Flow
FFJORD
396K
396K
5.43e-1
5,8001
6.76e-1
6.4%6-1
-104.9
-37.80
IONN
304K
6.74-1
1446-2
-1.65
IGNN
990K
5.54e-1
6.30c
-37.68
IRoNNet
304K
3.20e-1
2.736-2
-1.17
IResNet
990K
5,50e
5.50e
-33.11
ScoreSDB
76K
1.05e
8.26e
3.89
ScoreSDE
396K
5.63c-1
6.00e1
-7.55
time. The performance is also reported under a fixed-budget setting to compare across models. The numerical inversion error of the trained flow models are at the le-5 level or below, see Table
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 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 JKO-iFlow gives satisfactory generative performance. The results of JKO-iFlow in Figure are obtained after applying the trajectory improvement techniques in Section
4.2
The comparison results without using the techniques are provided in Figures A.2 and A.3 which
show the benefit of the two techniques in learning the model.
5.3
High-dimensional tabular data
Table assesses the performance of JKO-iFlow and baseline methods four high-dimensional tabular datasets under a fixed-budget setting: we keep the model size of continuous flow models (OT-Flow and FFJORD) and the diffusion model (ScoreSDE) the same, and increase model sizes for discrete flow models (IGNN and IResNet) which can be computed faster therein and this also improves their performances. Except for MINIBOONE dataset, JKO-iFlow yields smaller or similar quantitative metrics compared to all alternatives. The visual comparisons of generated data distributions are shown in Figure A.44
After additional training to apply the trajectory reparametrization, the performance of JKO-iFlow is furtherly improved, see Table A.3 in comparison with additional baselines of OT-Flow and FFJORD from the original references. The effects of trajectory reparameterization on the MINIBOONE dataset shown in Figure [A.5 where the W2 movements across the blocks become more as the iteration proceeds, see (a). The generative performance, as shown in (b), visually improves after the reparametrization, which is consistent with the lower NLL values in Table A.3 (NLL improves from 12.55 to 10.55). These results show that JKO-iFlow performs comparably to the best baseline, and often better under small computational and memory budgets.
5.4 Image generation
We apply the JKO-iFlow model to image generation tasks on the MNIST, CIFAR-10, and Imagenet- 32 datasets. In all examples, we train JKO-iFlow in the latent space of a pre-trained variational auto-encoder ( (VAE) adopted from [Esser et al.) 2021]. Training generative models in latent space statc-of-the-art the-art image generation performance, c.g., in StablcDiffusion RRombach et al., 2022]. Here we train a flow model instead of score-based diffusion model in the latent space, and more details can be found in Appendix C.2
have been shown to obtain state-of-
The generated images are shown in Figure The quantitative evaluation is by the Fréchet inception distance (FID) score [Heusel et al., 2017). which is included in the figure captions. Comparatively, we are aware of some performance gap between our model and the state-of-the
performance of score-based diffusion models like DDPM (Ho et al., 2020) and ScoreSDE (Song et al.. 2021], however, the results here are obtained with less computation. Specifically, on a single 1 GPU, our experiments took 24 hours on CIFAR 10 and 30 hours on Imagenet-32. Meanwhile, the image generation by JKO-iFlow model here obtains visually more appealing images and achieves lower

---

0
8
6
4
3
3
4
9
4
8
8
S
3
s
/
O
5
7
0
9
9
0
3
9
02800
{
7
4
3
5
7
7
2
(a) Generated MNIST digits. FID: 7.95.
(b) Generated CIFARIO images. FID: 29.10.
(c) Generated Imagenet-32 images. FID: 20.10.
Figure Genernted samples MNIST,
Imagenet-32 CIFAR1O and I image per class for Imagenet-32. The FIDs are shown in subcaptions. Uncurated samples are shown in Figure A.6
model space. select images per
FIDs compared most CNF baselines Finlay et al. 2020
Grathwohl al., 2019 Behrmann et al.] 2019 Chen al.
2019
5.5
Conditional generation
The problem aims to generate input samples X given a label Y from the conditional distribution XJY to be learned from data. We follow the approach in IGNN [Xu et al. 2022]. In this setting, the JKO-iFlow network pushes from the distribution X = k to the class-specific component in the Gaussian mixture of Y k. see Figure A.7b and Appendix CC2.4 for more details. We apply JKO-iFlow to the Solar ramping dataset and compare it with the original IGNN model, and both models use graph neural network layers in the residual blocks. The results are shown in Figure AA where both the NLL and MMD-m metrics indicate the superior performance of JKO-iFlow and is consistent with the visual comparison.
6
Discussion
The work can be extended in several directions. The application to larger-scale image datasets and larger graphs will enlarge the scope of usage. To overcome the computational challenge faced by neural ODE models for high dimensional input, e.g., images of higher resolution, one would need to improve the training efficiency of the backpropagation neural ODE in addition to the dimension reduction techniques by VAE been explored here. Another possibility is to combine the JKO-iFlow scheme with other backbone flow models that are more suitable the specific tasks. Meanwhile, it would be interesting extend the method to other problems for which CNF models have proven to be effective. Examples include multi-dimensional probabilistic regression [Chen et al. 2018), plug-in deep architectures such as StyleFlow [Abdl et al. 2021 and the application to Mcan-field Games (Huang et al. 2023].
Theoretically. the expressiveness of the flow model to generate regular data distribution can be analyzed based on Section 3.3 To sketch road map, a block-wise approximation guarantee of f(x, t) as in (12) can lead to approximation of the Fokker-Planck flow (3). which pushes forward the density to be e-close to normality
= (1e)) time, see (4). Reversing the time of the ODE then leads to an approximation of the initial density Pn px by flowing backward in time from T to zero. Further analysis under technical assumptions left to future work. During the time that this paper was being published, the convergence analysis of the proposed model studied in Cheng al. (023)
10

---

Acknowledgement
The authors thank Jianfeng Lu, Yulong Lu, and Yiping Lu for helpful discussions. The work is partially supported by NSF DMS-2134033 C.X. and Y.X. are partially supported by an NSF CAREER CCF-16509) and NSF DMS-2134037, CMMI-2015787, CMMI-2112533, DMS-I993 1938106, and DMS- 183--1832200 and the Coca-Cola Foundation. XC is also partially supported by NSF DMS-223786 and Simons Foundation.
References
Rameen Abdal, Peihao Zhu, Niloy Mitra, and Peter Wonka. Styleflow: Atrrbute-conndiione exploration of syyllgan-generaaue images using conditional continuous normalizing flows. ACM Transactions on Graphies (ToG), 1023) 2021.
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. Optimizing functionals on the space of probabilities with input convex neural networks. Trunsactions on Machine Learning Research, 2022. ISSN 2835-8856.
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. International Conference on Machine Learning, pages 146-155. PMLR, 2017.
Miguel A Arcones and Evarist Gine. On the bootstrap of and statistics. The Annals of Statistics, pages 655-674, 1992.
Bolton Bailey and Matus Telgarsky. Size-noise tradeoffs 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):330-945, 1993.
Jens Behrmann, Will Grathwohl, Ricky TQ Chen, David Duvenaud, and Jörn-Henrik Jacobsen. Invertible residual networks. In International Conference Machine Learning, pages 573-582. PMLR, 2019.
Adam Block, Youssef Mroueh, and Alexander Rakhlin. Generative modeling with denoising auto- encoders and langevin sampling. arXiv preprint arXiv:2002.00107 2020.
Nicholas M Boffi and Eric Vanden-Fijnden. Probability flow solution of the fokker-planee equation. Machine Learning: Science and Technology, 4(3):035012, 2023.
François Bolley, Ivan Gentil, and Arnau Guillin. Convergence to equilibrium in wasserstein distance for fokker-planck equations. Journal of Functional Analysis, 263(8):440--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, pages
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ömn-eerik Jacobsen. Residual flows for invertible generative modeling. Advances in Neural Information Processing Systems, 32, 2019.
Xiuyuan Cheng, Jianfeng Lu, Yixin Tan, and Yao Xie. Convergence of flow-based generative models via proximal gradient descent in wasserstein space. arXiv preprint arXiv:2310.1753 2023.
Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. Advances neural information processing systems, 29, 2016.
11

---

Pierre Degond and S Vallic. The weighted particle method for convection-diffusii equations. i. the case of an isotropic viscosity. Mathematics of computation, 53(188):485-507, 1989.
Pierre Degond and Francisco-José Mustieles. A deterministic approximation of diffusion equations using particles. SIAM Journal on Scientific and Statistical Computing, 11(2:2993-31, 1990.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 JEEE conference on computer vision and pattern recognition, pages 248-255. leee, 2009.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6): 141-142, 2012.
Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: non-linear independent components estimation. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Workshop Track Proceedings, 2015.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high- righuressoution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pages 12873-12883, 2021.
Jiaojiao Fan, Qinsheng Zhang, Amirhossein Taghvaei, and Yongxin Chen. Variational Wasserstein gradient flow. In Proceedings of the 39th International Conference on Muchine Learning, pages 1115-6215. PMLR, 2022.
Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
Chris Finlay, Jörn-Herrik 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.
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.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, and David Duvenaud. Scalable reversible generative models with free-form continuous dynamics. In International Conference on Learning Representations, 2019,
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, 2012a.
Arthur Gretton, Dino Sejdinovic, Heiko Strathmann, Sivaraman Balakrishnan, Massimiliano Pontil, Kenji Fukumizu, and Bharath K Sriperumbudur. Optimal kernel choice for sarge-scale two-sample tests. Advances in neural information processing systems, 25, 2012b.
Ishaan Gulrajani, Faruk Ahmed, Martín Arjovsky, Vincent Dumoulin, and Aaron C. Courville. Improved training of wasserstein gans. In NIPS, 2017.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Berhaard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Chin-Wer Wei Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective on diffusion-based generative models and score matching. Advances Neural Information Processing Systems, 34: 2021.
12

---

Han Huang, Jiajia Yu, Jie Chen, and Rongjie Lai. Bridging mean-field games and normalizing flows with trajectory regularization. Journal of Computational Physics, page 112155, 2023.
Michael F. Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Stafistics Simulation and Computation, 18:1059-1076, 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), pages 5967-5976, 2017.
Rie Johnson and Tong Zhang. framework composite functional gradient methods for generative adversarial models. IEEE transactions pallern analysis and machine intelligence, 43(1):17-32, 2019.
Richard Jordan, David Kinderlehrer. and Felix Otto. The variational formulation of the fokker planer equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
Diederik P. Kingma and Jimmy Ba. Adam: method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun, editors, 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
Diederik Kingma and Max Welling. An introduction variational autoencoders. Foundations and Trends@ in Machine Learning, 12(4):307-392, 2019.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems, 31, 2018.
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.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference Learning Representations, 2021.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical Report University of Toronto, Toronto, Ontario, 2009.
Holden Lee, Rong Ge, Tengyu Ma, Andrej Risteski, and Sanjeev Arora. On the ability of neural nets to express distributions. In Conference on Learning Theorr, pages 1271-1299 PMLR, 2017.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Repre- sentations, 2023.
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 Robban Stöter.
Sliced-wassersteii flows: Nonparametric generative modeling via optimal transport and diffusions. In International Conference on Machine Learning, pages 4104-4113, PMLR, 2019.
Yulong Lu and Jianfeng Lu. A universal approximation theorem of deep neural networks for expressing probability distributions. Advances neural information processing systems, 33: 3094-3105, 2020.
13

---

James Lucas, G. Tucker, Roger B. Grosse, and Mohammad Norouzi. Understanding posterior collapse in generative latent variable models. In DGS@ ICLR, 2019.
Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions fokker-
planck equations through gradient-log-denstty estimation. Entropy, 22(8):802, 2020.
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, 012115243-55555 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 neural information processing systems, 30, 2017.
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 Vetito, 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é-Bor Buc. E. Fox, and R. Garnett, editors, Advances Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019.
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.
Dmytro Perekrestenko, Léandre Eberhard, and Helmut Böleskei. High-dimensional distribution generation through deep neural networks. Partial Differential Equations and Applications, 2(5): 1-44, 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pages 10684-10692 2022.
Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2022.
Lars Ruthotto, Stanley J Osher, Wuchen Li, Levon Nurbekyan, and Samy Wu Fung. A machine learning framework for solving high-dimensional mean field game and mean field control problems. Proceedings of the National Academy of Sciences, 1 11117):9183-919 2020.
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, Xi Chen, and Xi Chen. Improved techniques for training gans. In Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
Antonin Schrab, Imun Kim, MÅ@lisande Albert, BÂ@atrice Laurent, Benjamin Guedj, and Arthur Gretton. Mmd aggregated two-sample test. Journal of Machine Learning Research, 24(194)1-88. 2023.
Zebang Shen, Zhenfu Wang, Satyen Kale, Alejandro Ribeiro, Amin Karbasi, and Hamed Hassani. Self-consistency of the fokker planck equation. In Conference on Learning Theory, pages 817-841. PMLR, 2022.
Thomas Sideris. Ordinary differential equations and dynamical systems, volume 2. Springer, 2013.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In International Conference on Learning Representations, 2015.
14

---

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 on Learning Representations, 2021.
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.
Danica Sutherland, Hsiao-Yu Tung, Heiko Strathmann, Soumyajit De, Aaditya Ramdas, Alex Smola, and Arthur Gretton. Generative models and model criticism via optimized maximum mean discrepancy. In International Conference on Learning Representations, 2017.
Esteban Tabak and Eric Vanden-Eijnden. Density estimation by dual ascent of the log- lieliikelioood. 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 a90iv:1905. 008833, 2019a.
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference Learning Theory, pages 3084-3114, PMLR, 2019b.
Alexander Vidal, Samy Wu Fung, Luis Tenorio, Stanley Osher, and Levon Nurbekyan. Taming hyperparameter tuning in continuous normalizing flows using the jko scheme. Scientifie Reports, 13(1):4501, 2023.
Chen Xu, Xiuyuan Cheng, and Yao Xie. Invertible neural networks for graph prediction. IEEE Journal on Selected Areas in Information Theory, 3(3):54-4677 2022.
Dmitry Yarotsky. Error bounds for approximations with deep relu networks. Neural Networks, 94: 103- 2017.
Linfeng Zhang, Lei Wang, al. Monge-Ampére flow for generative modeling. arXiv preprint 17X1v:180,.10188, 2018.
Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in Neural Information Processing Systems, 15116280-1229,, 2021.
15

---

A Proofs
A.1 Proofs in Section
Lemma A.I. Suppose and are two densities on E in P, the following two problems
min Lplp] (II(pllg) PEP
(2 (p.p).
(13)
2h
1
min Lr(T]= = - + Eguple -(x))?
(14)
Tike -2
2h
have the same minimum and (a) T TR Rdd is minimizer of (14), then p² )T
# is minimizer of
13
(b) If p² a minimizer of (13). then the optimal transport from p to p² minimizes 14
Proof of Lemma A.I Let the minimum of 14 be Li,, and that of (33 be Lip-
Proof of (a): Suppose LT achieves minimum T then T the optimal transport from p to
p² (T*) #p because otherwise Lt can be further improved. By definition of T.p. we have L. We claim that L 1. Otherwise, there another p' such that Lplp"l < Li. Let T be the optimal transport from p to p', and then L1 [T") Lplp' Li-
Tij [T(
T,pl="]
contradicting with that L is the minimum of Ly. This also shows that L = L = Li that is, p² a minimizer of Lpp
Proof of (b): Suppose Lpp achieves minimum at p" Let 7 be the OT from p to p², then Ezzp// T (2) = W2(p, )² and then Lr[T* = L,(e"] LG which equals L7 as proved in (a). This shows that T is minimizer of Lr-
Proof of Proposition 3.1
Given pk being the density of (tt) at t = kh, recall that T is the solution map from s(()) s h). We denote Pt Pk. and P++h = Pk. By definition,
KL(T@Pellpz) = Ex-peth (log Pe+h (x) - log Pz)))).
(15)
Because pz (zz) //, we have logpz Vv + c₁ for some constant c1 Thus
Ezeplen log p2 = (t(t) log P2 (z(t h)) = c1 Em(t)ra V(r(t + h)).
(16)
To compute the first term in 15 note that
E Ex-pet log (x+)(r) = Ex(()~p log PEth (x(t h)),
(17)
and by the expression (called "instantancous change-of-variable formula" normalizing flow litera- ture [Chen et al. 2018), which we derive directly below)
dt log p(x(t), = - f(x((),t),
(18)
we have that for each value of xzt).
++h V f()(s),
log Pe+- (x(t + h)) log p(x((= h), h) log (()()))
Inserting back to we have
r+h
E Excpees log Pt+h = Ex(t))p, log pr(z(t)) - Ex(t)=p
V - f(z(s)
The first term is determined by Pt = Pv. and thus is a constant 02 independent from f(x, t) on t [kh, (k 1)h]. Together with 16), have shown that
r.h.s. of 15p C2 Ex(t)ros
((o(s), - ei + Ex(t)~p V(a(t + h)),
which proves
7
16

---

Derivation of (18): by chain rule,
d
log p(x((),))) = Vp(æ(t),t) i(t) + a,p(a(t).t)
di
p(x(1),))
Vp. V (pf)
(by and 2
(=(0).)
- t()
B Technical details of Section
To train the proposed JKO-iFlow model by Algorithm 1 the k-th JKO block integrates from 1 to lx, starting from to = O. Denote by /1 := 4x+1 4 the step-size for k-th step.
one needs specify sequence of tus where
B.1 Initial choice of t in Algorithm 1
We first initialize the choice by the following scheme
h& = min(p ho, hmax) 1: = 0, 0,
(19)
where the base stepsize hgg the multiplying factor p 1, and the maximum stepsize Mmax > ho are three
to be specified by the user. When p 1, the sequence of h4 in (19
corresponds to constant stepsize h4 ho.
Applying Algorithm] [ with the initial choice l gives trained residual blocks 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 hg and we explain in detail below.
B.2 Hutchinson trace via finite difference
We propose Enite-difference estimator of V f(x,t) Tr(J+(z)): where JF(x)) is the Jacobian of f at x. This is based on the Hutchinson trace estimator [Hutchinson] 1989.. which states that Tr(Je(x)) = Ep(e) [e='J((x)e Here, p(e) is a distribution IR satisfying E(e] = and Cov(e) (e.g., standard Gaussian). Existing CNF works estimate the expectation by sampling e p(e) and computing the vector-Jacobial product Jr(I)e using reverse-mode automatic differentiation. The product Jr(z)c can be computed for approximately the same cost as evaluating [Grathwohl et al. 2019]
Now, given fixed we have (x(x) = lim f(a+te)- f(m) which is the directional derivative of f along the direction e. Thus, given fixed small Ö, we propose the following finite-difference estimator of V -f
I f(x + σe) f(x)
f(x,t) no Ep(e)
(20)
The finite-difference approximation of (r(z) requires only one additional function evaluation of at + and avoids the computation of automatic differentiation. We empirically found that training with (20) on high- nigh-dimensional examples (e.g., d > 100 as image examples) can be approximately 1.5~2 times faster than existing CNF approaches relying reverse-mode automatic differentiation. In all experiments, we let 0 Vd with do 0.02.
B.3 Trajectory improvement illustrated in vector space
introduces two techniques to improve the computation gradient descent trajectory represented by discrete points. For illustrative purposes, we introduce the two techniques when the gradient system in a vector space equipped with Fuclidean metrie. The JKO-iFlow training applies the same technique in the gradient system equipped with the Wasserstein-2 2 metric.
Section 4.2
17

---

Let the space be R" and F R + R be differential landscape. Given sequence of positive stepsize fig, one can compute a sequence of representative points Fy which discretizes the gradient descent trajectory of minimizing F (towards a local minimum x") Specifically, we have
3k+1 arg min F(x)
|| xx
(21)
2h
starting from some fixed ro-
B.3.1 Trajectory reparameterization
Starting from an initial sequence of step-size (-a)f-1 in Section B.I the reparameterization technique implements an iterative scheme to update the sequence hi adaptively. call j-th iteration 'Iter-j', and denote the sequence by (ht K The corresponding solution points
by solving (21) with h(") are denoted as x addition to the constant Amax. we need another algorithmic parameter (0, 1), and set 0.3 for the vector-space example in computation.
We always have zn zo for all 'Iter-j". In the j-th iteration of the parametrization, we compute the following
1. Compute the arclength as
(
10 (
=222122, h 1,..,
(22)
2. Compute the average arclength
$ :
((i) /T.
k=
3. Update the stepsize for k = 1, L as
[++1) min (1t (Sh
(1) (5) -1())) Nmaxx
hk
4. Solve (2 using A (i+1)
to obtain < (++1)
We terminate the iteration when the arclength (5(() are approximately equal. An illustration is given as Iter- 12 in the upper panel of Figure [. After the reparameterization iterations, we solve an additional F4.+1 by minimizing F(r) starting from ZT. This corresponds to optimizing the "free-block" in Algorithm
(
B.3.2 Progressive refinement
The scheme can be computed using positive integer refinement factor. For simplicity, use factor 2 throughout this work.
Given sequence of ( and the representative points {ax k the refinement scheme computes a refined trajectory having k ..... 21:
1. Compute the refined stepsizes
h2zk higk h2/2, k k= 1...., L
2. Compute the representative points 3y by solving 21 using he, possibly warm-starting by Σ2x T&k and T2k-1 = (It + 12/2
We then apply the trajectory reparameterization iterations in Section B.3. to the refined trajectory till convergence, and the free-block ending point also recomputed. An illustration given as r-Iter- 10 in the upper panel of Figure A.I
18

---

B.4 Trajectory improvement in probability space
We apply the techniques to solve the JKO step (8) which is the counterpart of (21) gradient descent scheme with proximal steps. The representative points Ią are replaced with transported distributions pk: which form sequence points in and the optimization is over the parametrization of induced by the neural network flow mapping consisting of the first & residual blocks.
It remains to define the arclength in 22 to implement the two techniques. Because equip P with the Wasserstein-2 metric, we compute (omitting the iteration index j in the notation)
S W2(Pk-1.P6) = (Exxpe- I (1)
(23)
where the transport mapping (x(x) = fhL fo. (r(s), s)ds and can be computed from the k-th is computed by finite-sample average the block. In practice, the expectation Ex-pu training set.
At last, the optimal warm-start of pa in the refinement is implemented by inheriting the parameters 0g of the trained blocks.
C Experimental details
C.1 Quantitative evaluation metrics
Besides visual comparison, we adopt two quantitative metries to evaluate the performance of gener- ative models, the negative log-likelihood (NLL) metric Grathwohl et al. (2019) and the maximum mean discrepancy (MMD) [Gretton et al. 2012a metric.
arc lengths mean std
initial:0.344 0.097
8te.-8:0.357 0.018
1tter3:0.184 0.009
Iter1:0.347 0.078
Iter-9:0.357 0.015
Mtter-4:0188 0.007
iter:2:0.350 0.062
Iter-10:0.357 0.013
riter-5:0.184 0.006
Iter:3:0.352 0.051
Iter-11:0.333 0.010
Aiter-6:01844 0.004
Iter-400.353 & 0.041
Iter-12:0.358 0.008
r-tter 7:0.184 0.003
iter:5:0.355 0.034
refine:0.183 & 0,023
riter-8:0.184 & 0.003
iter6:0.3555 0.027
-iter-1:0.189 0.017
riter-9:0.184 0.002
Iter:7:0.356 0.022
rilter 1:0e1-2:0.184 0.012
riteer10:001844 & 0.002
0.5
0.4
0.3
0.2
0.00
0.05
0.10
0.15
0.20
0.25
0.30
0.35
0.40
0.45
time Σh
Computed trajectory with L-16
3.0
2.5
2.0
1.5
1.0
0.5
0.0
16:
minumum: (-1.911,
Figure upper panel shows lengths and mean standard deviation lengths reparameterizattoo iterations,
one refinement, and an additional 10 reparameterization iterations. The lower panel visualizes the trnjectory consisting of 17 solution points "r-Iter- 10", where the free point 17-th point computes the approximate minimum 317 #241 = [-1.911, 0.105,
19

---

C.1.1 NLL metric
Our computation of the NLL metric follows the standard procedure for neural ODE normalizing flow models Chen al. [0118] Grathwohl et al. (2019), where the evolution of density p(x(tt) can be computed via integrating V [[x((),t) over time due to the instantaneous change-of-variable formula (8)).
Specifically, our model flows from 4g = 0 to LL+1 T, where xr()) ~ Po the data distribution, and we train the flow model to make x(T') follow a normal density q. The model density at data sample expressed as p(z, 0) (I,5) ) ((x) where 77 is the inverse model flow mapping from x(T') to x <0)). Thus the model log-likelihood can be expressed as
/ V. f(I(s),
log ((x((), 0) = log
In practice, our trained flow model has L blocks where each block fe, represents f(r, t) on (ts-1,tx) and is parametrized by 0z. The log-likelihood test sample is then computed by
1 LL(x)= - -11 dll(:==))))(++=o(())))))))))))))
Σ
L-1
v (o,)
(24)
where
x(tx) x(te-1) + ft for (x(s),
starting from x00) x. Both the integration of fo, and V fo. 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 X == {a; | and X = { ارار and a kernel function
k(r, the (squared) kernel MMD |Grettn et al. 2012a is defined as
1
M
2 NM
MMD(X, x)
,,), + M
k(),
k(r,, (25)
N2
=1
When generative model is trained, generate M i.i.d. data samples by the model to construct the set X, and we form the set X using N true data samples (from the test set). MMD metrics with other choices of kernels are possible Gretton et al. [2012b], SStherland et al. (017)) Schrrb et al.
[2023]. In all experiments here, we the Gaussian kernel k<(,,)) exp(- x कl /2h²) } stay consistent with reported baselines from [Onken et al. 2021]] where > 0 is the bandwidth parameter. We use three ways of setting the bandwidth parameter h:
Constant bandwidth: h= he = 1. The resulting MMD denoted as 'MMD-1'.
Median bandwidth [Gretton al. 2012a]]: let htm be the median of || - x, 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 a more powerful MMD test distinguish the difference in the underlying distributions) while ensuring that we use large enough M and N compute the MMD metric. We call the metric 'MMD-c'.
On all datasets, we use at most N = 10K test samples X, and for each trained model, we generate M = 10K test samples to form the dataset X to compute the MMD value defined as in (25) Note that the MMD metric as a measure of distance between two distributions significant when above test threshold which is defined as the upper (1 -u-uutle of the distribution of the MMD
20

---

HMD. 2.55
MHE 3.05U
(a) Per-block W2 over reparameterization iterations.
training (middle) and (right). NLL values shown the title.
Figure A.2: Same plots in Figure A.5 for Rose data. We observe improved generative quality the leaves of the rose after reparameterization
WiioL 1o0d
NHD Threshale
MMD:: M.L: 2.29 MHE 2.17c
Inttal
Ite
Itor
refine
Filer
her
Iter
[a] Par-block W22 over reparamcterization iterations and refinement
(b) Results (middle)) r-ter (righ). MMD NLL values are shown in the titl..
regarameterization lieration retinement).
Figure A.3: Same plots Figure A.5 for Fractal data. We observe improved generative quality the edges of the after refinement.
statistic under the null hypothesis (i.e., when dataset X and X observe the same distribution). The sealar is the controlled Type-I error (known the test level), which is set be 0.05. To obtain the test threshold for the MMD, we adopt the bootstrap procedure (Arcones and Gine 1992]. Gretton et al. 2012a] and compute r as the empirical (1 a))-uantile of the simulated null distribution of the MMD from the pool of samples formed by the union of X and X, where the set X is generated our usage
by the trained JKO-iFlow model. We use 1000 times of bootstrap in all experiments. baseline of the MMD metric, that is, when the computed MMD values are above then the smaller the MMD value the better the generative performance; when the computed MMD value below т, means that with respect to the current MMD metric, the trained model generates a data distribution that is as good as the true distribution.
of evaluating generative models, the threshold チ can be viewed
C.2 Detail setup and additional results
All experiments conducted using PyTorch Paszke et al. 2019 and PyTorch Geometric (Fey and Lenssen 2019). The optimizer Adam Kingma and Ba 2015 with learning rates to specified in each example. We the neural-ODE integrator [Grathwohl et al. 2019 adjoint method] all examples.
C.2.1 Two-dimensional toy data
Training and test data:
We generate test samples for each dataset, and for the training split,
For Rose (Figure and A.2) and Checkerboard and Olympic rings (Figure 3 re-sample 10K training samples every epoch, for total of 100 epochs per block. The batch size 500.
For Fractal tree (Figure 3 and A.3), we use fixed training data of 200K samples, for a total of 50 epochs per block. The batch size 2000.
Choice of lu: The initial schedule is by setting ho 0.75, 1.2. We set /max 5 for rose, checkerboard, and Olympic rings, and hmax for fractal tree. For reparametrization and refinement, we use 77 = 0.5 in the reparameterization iterations for all examples.
Table A.I: Inversion emror Ea-wy | (a)) 2 of JKO-iPlow computed via sample average on the test split of the data where
Te denotes the transport mapping over all blocks trained flow network.
Rose Fractul tree Olympic rings Checkerboard
POWER
GAS
MINIBOONE
BSD300
MNIST 1.87e55
3.30c-6
3.58c-5
2.24c-6
3.07e-5
1.8c-55
1.58c-6
1.09c-6
1.53c-5
21

---

Table A.2: Complete results real tabular datasets to augment Table under the fixed budget setting. denotes to the number of residual blocks used in each dataset; for comparison, the free block in JKO-iHlow is not used. For fäir compurison across models, the number of batches indicates how many batches pass through all residual blocks.
Data Set
Model
Param
Training
Testing
Time (h)
Batches
Tamee/Btcchescco (s)
Batch size
MMD m
MMD-1
NLL
JKo.aflow 76K, L-4
0.07
0.76K
3.51e-1
10000
9.86c-5
2.40c-4
-0.12
OT-FFlow
76K
0.36
7.58K
1.71e-1
10000
7.58c-4
5.35c-4
0.32
POWER
FFJORD
76K,
10000
IGNN
304K, L-16
0.29
7,58K
1.38c-
10000
1.93e33
1.59e-3
0.95
IReSNet
304K, L=16
0.41
7.5SK
1995-1
10000
3.92e-3
2.13-22
3.37
ScoreSDE
76K
0.06
7,58K
2.85e-2
10000
9.12c-4
6.08e-3
3.41
ScoreSDX
76K
75.8K
2.85e-2
10000
7.12e-4
5,04-3
JKO-FFlow
57K, L=3
0.05
0.76K
2.63c-1
10000
3.86e-4
7,20c-4
-0.06
L88e-4
1.73e4
JKD-HHlow
76K.
0.07
0.76K
3.33c-
5000
1.56-4
5,00c-4
-7.65
OT-Flow
76K
0.23
7.60K
1.09-1
5000
1.99c4
5.16c-1
-6.04
GAS
FFIORD
76K, 76K,
0.65
7.60K
3.08e-1
5000
1.87e3
3.28e-3
-2.65
IGNN
804K,
7.6IK
1.61e-1
5000
6.746-3
1.43-22
-1.65
IResNet
304K, L=16
0.46
7.60K
2.18c-1
5000
3.20c-3
2.73--2
-1.17
ScoreSDE
76K
0,03
7,60K
1,42e-2
5000
1.05c-3
8.36c-4
-3,69
ScoreSDE
2.23e-4
JKO-Flow
95K,
0.09
0.76K
4.15e-1
5000
1.51e-
3.77e-4
-7.80
2,46e-1
3.75e4
JKD-HHow
112K, Lm4
0.34K
3.61e-1
2000
6666-4
3.7Me-4
12.55
OT--Flow
112K
0.21
3.39K
2.23e-1
2000
6.58e
3.79e-4
11:44
MINIBOONE
FFJORD
112K, L=4
0.28
3.39K
2.97c-1
2000
3.51e-3
4.13c-4
23.77
IGNN
48K,
6.69p-1
IResNet
448K, L-16
0.71
3.39K
7.54e-
2000
2.13e-3
4.16c-4
22.36
ScoreSDg
112
0,01
3,39K
6,37e-3
2000
5. 5.86e-1
4.33c-4
27. 38
ScnreDDE
112K
33,90K
6.37e-3
2000
4.17e-3
3.87c-4
20.70
7: 1.38
T: 1.01c4
JKO-HHlw
996K, Lm4
0.05
1,3KK
1,85c-1
1000
2.24c-4
1.91e44
-153,82
OT-Fluw
BSDS300 63
FFIORD
396K, L-d
0.54
10.29K
1.89c=1
1000
5.60e-1
6.76e-
33.80
IGNN
990K, L-10
1.71
10.29K
5.98e-1
1000
5.64c-1
6.86e-1
-37.68
L=10
ScoreSDIE
10,29K
3,50k-3
1000
5.61-1
6.0011
-7.55
ScoreSDE
396K
0.10
102.90K
3.50--
1000
5.61e-1
6.62c-1
-7,31
JKO-Flow
396K,
0.08
1,03K
2.76e-1
5000
1.41e.4
8.83e-5
-156.68
For Rose, 3 reparameterization moving iteration, no refinement.
For Checkerboard and Olympic rings, 4 reparameterization moving iteration, no refinement.
For Fractal tree, reparameterization moving iteration, one refinement, and an additional reparameterization moving iteration.
MMD metric: We use custom kernel bandwidth h = 0.1 where hm is the median distance of dataset X with N 8K (from true data distribution). From trained generative model, we generate M 10K test samples as X.
Network and activation: Fully-connected residual blocks with two hidden layers. We use the softplus activation (3 = 20) with 128 hidden dimensions. Before refinement, we train residual blocks for rose, checkerboard, and Olympic rings, and train 6 residual blocks for fractal trees. The learning rate 5e-3.
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 the smallest edges are clearer using these techniques. In terms of quantitative metrics, the NLL also smaller, even though the MMD metric indicates there is no statistical difference between the ground truth and generated samples.
C.2.2 Tabular datasets
Training and test data: The four nigh-dimensional real datasets (POWER, GAS, MINIBOONE. BSDS300) come from the University of California Irvine (UCI) machine learning data repository, and we follow the pre-procsssing procedures of (Papamakarios al., 2017). Regarding data sizes,
POWER: 1.85M training sample and 205K test sample.
GAS: IM training sample, 100K test sample.
22

---

Table A.3: The OT Flow and FFFORD haselines marked with are from the original papers
Onken et al.
2021
and Grathwohl et al.
2019
The results of JKO-Fooww are obtained alter applying the reparamcterization technique.
Data Set
Model
# Param
Traminn
Tesing
Batches
Batch
NLL
POWER = 6
OT.Flow"
JKOHFoow
15K,
18K
6
22K
10000
10000
-0.30
-0.40
HHORDD
L
10000
GAS 8=8
1ow-Flow
HAK, LEb
127K
608K
52K
2000
5000
-9.20
-943
FHIORD"
279K. 1
-8
d 43
MINBBOONE
FT-Flow
FFJORD
KO-FFlow
S21K. L=1
7SK
T12K, L=1
1.72K
7K
2000
2000
1000
10.43
10.55
10.35
ESDS300 d E 63
O1HHoow
KO-FFoo
495K, L=5
207K
206K
37K
5000
300
-154.20
157.25
FFJORD
6.7M. L=2
10000
-157.10
MINIBOONE: 32K training sample, 3.6K test sample.
. BSDS300: 1.05M training sample, 250K test sample.
Choice of hg: The initial schedule is by setting hg 1. p 1, Amax 3. For reparametrization and refinement, we use = 0.5 in the reparameterization iterations for all datasets.
MMD metric: When the test data from the MINIBOONE dataset is less than 10K in size, we use all the test sets X and N the size of test data. When the test data from all three other datasets has more than 10K samples, we randomly select N = 10K samples from the entire test data, and the same random test subset is used to evaluate all models. We use the median distance kernel bandwidth and generate M- = 10K test samples from each model to form X.
Network and activation: We use the softplus activation for all networks. Regarding the design of residual blocks, we use fulne-con residual blocks with two hidden layers. On POWER, GAS, and MINIBOONE, we use 128 hidden nodes, and on BSDS300, we use 256 hidden nodes. The
number of residual blocks for four tabular datasets is described in Tables A.2 and A..3 Regarding learning rate, it is 1e-3 on POWER, 2e-3 on GAS, 5e-3 on MINIBOONE, and 1. .75e-3 on BSDS300.
Additional results: For experiments under the fixed-budge setting, Table A.2 shows the complete results of JKO-iFlow against competitors on tabular datasets. In the extra lines in the table, we show the result of JKO-iFlow with L=3 for POWER and L=5 for GAS, as these numbers of L are determined by the termination criterion Algorithm [] (The main lines show the results with so that our model has the same capacity as the alternatives.) BSDS300 has an extra line for JKO-iFlow trained with batch sizes 5000, which improves over the result with batch size 1000. For ScoreSDE. we use the implementation in (Huang et al., 2021]. which, during training, maximizes the evidence lower bound (ELBO) as an equivalent objective to the score-matching loss. Although ScoreSDE the fastest, its performance, even under 100 times more mini-batch stochastic gradient descent steps than JKO-iFlow, is still worse than JKO-iFlow in terms of both MMD-m and NLL.
To visualize the generative performance, scatter plots of the generated samples by JKO-iFlow and competitors on these tabular datasets are shown in Figure A.4 after projected to dimensions (determined by principal components computed from true data test samples). The plots show closer match between those from JKO-iFlow with the ground truth in distribution. For experiments allowing more expensive models and longer training time, Table A.3 shows the results of JKO-iFlow after additional training of trajectory reparametrization in comparison with other baselines cited from the original papers. We run 7 reparameterization iterations on POWER, GAS, and MINIBOONE and 1 reparameterization iteration on BSDS300. A free block used on POWER, GAS, and BSDS300.
Table A.4: NLL per noise scheduler and Buae combination of ScoreSDE on MINIBOONE The three settings "linear, constant, quadratic"
follow the DDPM suggestion [Io el 2020 mean standard deviation are computed over replicas of the trained model.
Noise scheduler
10
15
20
Linear
24.51 (0.24)
18.73 (0.64)
20.28 (0.38)
26.76(1.77)
26.33
Constant
25.47 (0.25)
30.37 (0.84)
37.73 (0.98)
40.47
45.32 (0.40)
Quadratic
27.19 (0)
17.45 (0.17)
18.48 (0.38)
18.90 (0.48)
21.35 (0.60)
23

---

Table A.S: Testing NLL noise scheduler and Smax combination of ScoreDDE on MINIBOONE, using larger network with longer training.
The table format is identical to that A.b
Noise scheduler Pmaxx
Linear
(0:44)
(L.77
10.47 (0.13)
Quadratic
12.88 (0:40)
11.11 (0.24)
11.05 (0.49)
True
Jo-iflow
#FjoRD
Se ScreSDE
(a) Power
True
|K0-FFlo
OT-Fipe
FFOORD
IGNN
ScoreSDE
Gas
True
KC iFlw
OT-Fiow
FFJORD
Past
Score5DE
(c) MINIBOONE
JKG-iflew
Flo
FIJORD
IGYN
ScoreSDE
(d) BSDS300
Figure A:4; Generative quality tabular datasets PCA projection of generated samples. The generative quality general aligns with the quantitative metrics in Table and
To ensure fair comparison against diffusion models, perform additional experiments using dif- ferent noise schedulers (t) and 8max in ScoreSDEE on MINIBOONE, following the noise scheduler suggestions in DDPM (Ho et al. 2020]. Note that DDPM be viewed discrete-time version of the variance- variance-preserving ScoreSDE model. As shown in Table A.4, the performance of ScoreSDE indeed sensitive the noise schedule. However, the best NLL of Score 17.45 from the table is still noticeably higher than the NLL of 12.55 obtained by JKO-iFlow on MINIBOONE in Table A.2, and JKO-iFlow trained using ten times less number of batches. To improve the performance of ScoreSDE on this example, we use ScoreSDE without restrictions on modeling and computation. Specifically, we consider larger network following the setup in [Albergo and Vanden Vanden-Eijnden 2023), where the network has 4 hidden layers with 512 hidden nodes per layer and consists of 831K parameters total. We then train this larger model for 100K batches using ScoreSDE with various noise schedulers, choosing noise schedulers that performed best based results in Table A.4 From the results reported in Table AA.5 we see that the testing NLL by ScoreSDE can be as low as 10.47, which is among the state-of-the- the values reported in Table A.3 Nevertheless, we want to highlight that the proposed JKO-iFlow can obtain competitive results (i.e., 10.55 in Table A.3) with much smaller models: each JKO-iFlow block has only hidden layers with 128 hidden nodes per layer, and we trained 4 blocks that contain 12K parameters total. We also trained JKO-iFlow for only 2.72K batches, rather than the 100K batches used for ScOreSDE.
C.2.3 Image generation with pre-trained variational auto-encoder
We perform image generation MNIST Deng 2012 CIFAR10 Krizhevsky and Hinton 2009 and Imagenet-32 [Deng et al. 2009 datasets. We do so in the latent space of pre-trained VAE, where we discuss the details below.
24

---

VAE as data pre-processing: We train deep VAEs in an adversarial manner following Esser et al. 2021] and use pre-trained VAES pre-process the input images X. Specifically, the encoder 8 (u(X), X(X)) of the VAE maps a RGB-image X to parameters of multivariate Gaussian distribution (x))), ((X) in lower dimension d. Then, given any random latent code X ()) ((((())>(())) the decoder D X > X of the VAE is trained that X ~ X for the reconstructed image X- ). On MNIST, we let the latent-space dimension d 20, and on CIFARIO and ImageNet32, we let the latent-space dimension = 192. There are 5.6M parameters in the VAE encoder and 8.25M parameters in the VAE decoder. Given trained VAE, we then train the JKO-iFlow model I' to transport invertibly between the distribution of random latent codes X defined over all training images X and the standard multivariate Gaussian distribution in IR
Training and test data: Training data: 60K images in MNIST, 50K images in CIFARIO 10, and 1.28M images in Imagenet-32. Test data: 10K images in MNIST and CIFARIO, and 50K images in Imagenet-32.
Choice of hg: on MNIST, we specify (hg = 1,p hmax = 1}. on CIFARIO, we specify (ho 1 1, hmax 1]. On Imagenet-32, we specify (hg 1, 1.1, Amax oo). We train L 6 blocks on MNIST, and we train = blocks on CIFAR1O and Imagenet-32. On CIFARIO and Imagenet-32, we also scale hgg = T hx/ Ej h3 so Ex ha T. We let T 0.8 on CIFARIO and T 0.4 on Imagenet-32.
Network architecture: We use the softplus activation with 3 = 20 for all hidden layers. On MNIST, use fully connected residual blocks with three hidden layers 256 hidden nodes. On CIFARIO and Imagenet-32, we parametrize each fo, as a concatenation of convolution layers and transposed convolution layers. Specifically:
CIFAR10: convolution layers have channels 11-644---220 100-256-255 256 with kernel size 3 and strides 1-1-2-1. Transposed convolution layers have channels 256-256-128-644 with kernel size 3-4-3-3 and strides 1 1-2-1-1. Total 2. 1M parameters.
Imagenet-32: convolution layers have channels 3- -2----4-122------2 strides 1-1-2-1. Transposed convolution layers have channels $12-256-128-64-3 with kernel size 3-4-3-3 and strides 1-2-1-1. Total 3.3M parameters.
with kernel size 3 and
We remark that because inputs into the JKO-iFlow blocks are latent-space codes that have much lower dimensions than the original image, our JKO-iFlow blocks are simpler and lighter in design than models in previous NeuralODE [Grathwohl et al. 2019 |Finlay et al. 2020] and diffusion model works [o et al.
2020 Song et al. 2021] Boffi and Vanden- Mijienn--indenn 2023 For instance, the DDPM model Ho et al.
2020 on CIFAR 10 has 35.7M parameters with more sophisticated model designs.
Training specifics; On MNIST, we the batch size to be 2000 during training, and we train 15K batches per JKO-iFlow block. We fix the learning rate to be 1c-3.
On CIFARIO and Imagenet-32, we fix the batch size to be 512 during training, and we train 75K batches per JKO-iFlow block. The time per batch is 0.18 seconds on Imagenet-32 and 0.15 seconds on CI.A10. The total time on Imagenet-32 30 hours and on CIFAR1O is 24 hours. When training the blocks, the initial learning rate 1r0 is decreased by a constant factor of 0.9 every 2500 batches.
Initi
Bar
iter
--- Bar
Bar
Perblock W reparameterization iterations.
(b) Results initial training (midlle) and Iter (right)..
Figure A.S: Reparametrization iterations JKO iFlow mudel on MINIBOONE, (a) reparameterizzaiion iteration,, a trajectory more uniform W2 movements is ohtiied. (b) Generated samples by the models hefore and after the moving iterations (visualized in the first two principal components computed from data).
25

---

Generated CIFARIO images
(b) Generated Imagenet-32 images
Figure Uncurated generated samples CIFARIO Imugenet- by FFlow in space.
XI- "(HIV)
HJY
(a) Two-cireles data (unconditional)
(b) Two moons data (conditional)
Figure A.7: Unconditional and conditional generation on simulated toy datasets by JKO-iFlow. color samples by class taking binary values. In both (a) and (b), generated samples close true samples in distribution. In (b), the pushforward distribution by JKO-iFlow, denoted as H|Y, also close to the target Gaussian mixture distribution HYY.
We let 1r0 be le-3 when training the first 5 blocks Imagenet-30 32 and decrease it to be 8c-4 on blocks 6-8. We let 1r0 be 1e-3 when training the first blocks on CIFAR 10 and decrease it to be 7.5c-4 on the rest 6 blocks. Additionally, we use gradient clipping of max norm 1 after collecting gradients on each mini-batch.
Additional results: We perform the following steps curate generated samples for CIFARIO and Imagenet-32, which are shown in Figure 4 we first train an image classifier (i.e., VGG-16 (Smonnyan and Zisserman, 2015)) on the training data. Then, we generate large number of uncurated images (200K for Imagenet-32 and 20K for CIFARIO), classify them using the pre-trained classifier, and sort them based top-1 predicted probability. Out of the top 20K images for Imagenet-32 and top 750 images for CIFAR 0 upon sorting, we manually select images that we think most resemble the given predicted class. Figure A.6 further shows uncurated images on CIFAR10 and Imagenet-32 by the same IKO-iFlow model after training.
C.2.4 Conditional generation
To modify the JKO-iFlow to apply to the conditional generation task, follow the framework Xu et al. [2022], which trains single flow mapping from X to H that pushes to match each component of the distribution associated with a distinct output label value = k., I.... K. Specifically, for data-label pair (X;, Y.). consider the continuous ODE trajectory æ(() starting from zx X
the per-sample training objective by changing the term V(=((++1) 1)) in ( to be Vy, (x((++))
where V (-) is the potential of the Gaussian mixture component H k. Because the Gaussian
26

---

(a) True X|Y 2.13--3, MMD-m:
IKO-Fllow
(b) IGNN
(c) True X
JKO. FFoo
(d) IGNN
8.12c-3
5.70e-2
2.88-3, MMD-m:
4.00--2
7.67p-2
NLL
-13.60
-4.55
NLL
4.19
1.53
Figure A.8: Conditional graph nodal data generation JKO-iFlow iGNN Solar ramping data. The nodal X R and the
nodal label Y E {0, Plots (altb) and (eXd) are for two different sets of values across all nodes respectively Samples X (concatenated by peojection
across nodes)
principal compunents (determined by true samples X
mixture is parametrized by mean vectors (the covariance is isotropic with fixed variance) [Xu et al. 2022
the expression of (+(() is quadratic function with explicit expression.
Training and test data:
For the simulated two-moon data, we re-sample 5K training samples every epoch, for total 40 epochs per block. The batch size 1000.
The solar dataset is retrieved from the National Solar Radiation Database (NSRDB), follow- ing the data procoesoosing in [Xu et al., 2022]. The dataset has a IK training samples and 1K test samples. The batch size 500, and each block is trained for 50 epochs.
Choice of hy: For both datasets, hg = 1, 1, 1mxx = 3. The reparamctrization and refinement techniques are not used.
MMD metric: The MMD values are computed for each specific value of Y (across all graph nodes). When is fixed, we retrieve test samples from true data, denoted as X YY, and generate model distribution samples, denoted as Y and then compute the MMD values the same as in the unconditional case. Due to the relatively small sample size of X (679 and 144 respectively), report the average values of MMD-m and threshold T over 50 replicas. In each replica, we subsample 90% observations of Y|Y and generate M = 2000 samples to form X|Y. For the results in Figure A.8 for the Y in plots (b), N 612 samples are randomly sampled from X|Y in each replica, and the standard deviation of MMD-m values is less than the 2e-3 level. For the Y in plots (c)(d), N 130 samples are randomly sampled from X|Y in each replica, and the standard deviation of MMD-m values is less than the 5e-3 level. The median distance kernel bandwidth computed on the entire X Y is used in computing the MMD values.
Network and activation: Regarding design of residual blocks,
On two-moon data, we use fully-connected residual blocks with two hidden layers under 128 hidden nodes. The activation is Tanh. We train four residual blocks. The learning rate 5e-3.
On solar data, we follow the same design of residual blocks as [u et al. 2022 More precisely, each residual block contains one Chebnet [Defferrard et al. 2016 layer with degree 3, followed by two fully-connected layers with 64 hidden nodes. The activation is ELU. We train 1l residual blocks. The learning is 5e-3.
27


Reviews:
[{"review_id": "ah8ACZ3rZC", "reviewer": "Reviewer_NZGt", "summary": "The authors introduce a novel normalizing flow training algorithm that integrates continuous normalizing flows and the JKO scheme. The objective of the training algorithm is to minimize KL divergence between the current density and the equilibrium density. The proposed method is inspired by the JKO scheme to enable adaptive blockwise training of residual networks. Furthermore, the authors demonstrate their method on synthetic and open-source datasets.", "questions": "1). Is it possible for the authors to provide FID scores for both MNIST and CIFAR-10?\n\n2). How good is the proposed method without using a pre-trained auto-encoder?\n\n", "limitations": "The limitations are clearly discussed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1). The proposed method is computationally efficient compared to other normalizing flow methods or diffusion. \n\n2). Novel idea of optimizing the velocity field instead of the transport map motivated JKO-Scheme. \n\n3). The paper is well-motivated and easy to understand. ", "weaknesses": "1). The paper lacks metrics/empirical validation on larger images. FID metrics for Figure 4 should be necessary considering the paper is about generative modeling and its challenging to tell how adequate the CIFAR-10 samples are because they have such low resolution via eye test.\n\n2). A significant highlight is that block-wise training leads to a large reduction in computation cost compared to other popular generative models (Figure 4), but the authors do not highlight that the proposed scheme uses a pre-trained auto-encoder, where some of the other methods do not. \n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Please refer to the global response for the questions on experiments of image generation and FID scores. \n\n* **(Weakness 2 & Question 2): The proposed scheme uses a pre-trained auto-encoder, where some of the other methods do not. How good is the proposed method without using a pre-trained auto-encoder?**\t\n\nWe thank the reviewer for asking the good question. We mainly propose to use the current method in the latent space for image generation tasks, which will be clarified in the revised draft. Directly applying the current approach to image pixel space is possible, but would incur high computational cost due to the likelihood-based training of the neural ODE backbone model, see more in the global response on \"the extension to larger images\". Consequently, the computational advantage of the block-wise training would be discounted.\n\nAt the same time, we would like to point out that the block-wise training of the JKO scheme aims to compute a sequence of transported distributions that get closer to the normal density progressively, and this may be combined with other backbone models that are more suitable for computation in image pixel space. One possibility is by adopting the flow models trained to match the desired velocity field [Albergo et al., 2023; Lipman et al., 2023]. We will clarify the limitations and discuss future directions in the discussion section.\n\nRef.\n\nAlbergo M S, Vanden-Eijnden E. Building normalizing flows with stochastic interpolants. ICLR 2023.\n\nLipman Y, Chen R T Q, Ben-Hamu H, et al. Flow matching for generative modeling. ICLR 2023."}, {"review_id": "R2qu7MFxAC", "reviewer": "Reviewer_uXKQ", "summary": "This paper proposes learning an invertible normalising flow using a neural ODE as a unique solution to the Fokker-Planck Equation (FPE) of the transport problem from the data distribution to the equilibrium solution. The solution of the FPE is obtained by using the JKO scheme by formulating it as a variational problem. The critical contribution of the paper is that each residual block of invertible neural ODE corresponds to a JKO step and the training objective can be computed from pushed data samples through the previous blocks. This leads to a block-wise procedure to train the JKO-iFlow model which is computationally more efficient than an end-to-end approach in the literature. The paper also proposes a way to determine the number of blocks adaptively by adaptively reparameterize the computed trajectory in the probability space with refinement to improve the model accuracy and the overall computational efficiency. Lastly, they empirical show that their method is competitive or better generative performance compared to existing flow and diffusion models on synthetic and real data, ", "questions": "Could there be probability trajectories that are particularly \"stiff\" much like how there are stiff ODE that require very small step sizes for the numerical ODE method to be stable and have the resulting solution be close to the actual trajectory?  The FPE does not only appear in image generation, in mean-field game theory the fokker-planck equation describes the dynamics of the aggregate distribution of agents can JKO-iflow work well on these other problems ? ", "limitations": "The authors adequately address the limitations of their work, and there are no obvious potential negative social impacts of their work. ", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper is novel, and the idea of using iterative block-wise training is a good one. This strategy can be applied to other works as well and I think is a good idea. The formulation of the loss function from the JKO-step is also very well-written and clear. One strength of the paper is that their proposed method is very flexible as they used neural ODE with sufficiently small step sizes in the residual blocks to invertibility making their method very general and doesn't have any restrictive architecture like is typical of neural of other invertible flow models. ", "weaknesses": "The experiments are on relatively small datasets. It would be instructive to test how the model performs on larger datasets like miniImagNet or even CelebA dataset. The MMD loss could try to use different kernels or even mixed kernels to evaluate the quality better of the generative models. When comparing JKO-iflow against DDPM, it might be valuable to discuss what noise scheduling for DDPM is being used for each dataset, as DDPM is sensitive to the noise schedule used. Thus, it might lead to an unfair comparison if DDPM are not properly tuned. The authors can also include inverse problems in their experiments as many one of the advantages of DDPM is the fact they can easily be adapted to tackle inverse problems as well. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Please refer to the global response for the questions on experiments of image generation. \n\n* **Different kernels in MMD to evaluate generative models**\n\nWe agree with the reviewer that the MMD loss metric depends on the choice of the kernel, and a variety of kernels may be adopted for evaluating the trained generative model. In our experiments, we adopt the Gaussian kernel to stay consistent with reported baselines from the literature [Onken et al., 2021] for the purpose of comparison. Being aware of the sensitivity of the kernel MMD to the choice of the kernel parameter, we performed MMD calculations with different bandwidth selections. In addition, to ensure that the comparison of MMD values is meaningful, we applied bootstrap to compute the threshold of the MMD test under null (when two distributions are the same), and all values to be compared are above the threshold. The details of the MMD tests, including bandwidth selection, are provided in Appendix C.1.2. Further evaluation of the generative models using more metrics would be an interesting direction for future work.\n\n* **Noise schedule for DDPM to ensure fair comparison**\n\nWe thank the reviewer for the suggestion, and have conducted additional experiments of different noise schedulers following DDPM [Ho et al., 2020] to ensure fair comparison. Below, we use three schedulers (i.e., the $\\beta(t)$ definition in [Song et al., 2021]) with five different choices for $\\bar{\\beta}_{\\max}$ in the variance-preserving ScoreSDE model, for which DDPM can be viewed as a discrete-time version. As shown in Table 1 below, the performance of DDPM is indeed sensitive to the noise schedule. However, the best NLL of DDPM 17.45 from the table is still noticeably higher than the NLL of 12.55 obtained by JKO-iFlow, and the latter is trained using ten times less number of batches. \n\nTable 1: NLL per noise scheduler and $\\bar{\\beta}_{\\max}$ combination, the three settings ‘linear, constant, quadratic’ follow the DDPM paper [Ho et al., 2020]. The mean and standard deviation are computed over three replicas of the trained model\n|Noise Scheduler\\\\ $\\bar{\\beta}_{\\max}$|1|5|10|15|20|\n|-|-|-|-|-|-|\n|linear|24.51 (0.24)|18.73 (0.64)|20.28 (0.38)|26.76 (1.77)|26.83 (1.23)|\n|constant|25.47 (0.25)|30.37 (0.84)|37.73 (0.98)|40.47 (1.27)|45.32 (0.40)|\n|quadratic|27.19 (0.14)|17.45 (0.17)|18.48 (0.38)|18.90 (0.48)|21.35 (0.60)|\n\nRef.\n\nHo et al. Denoising diffusion probabilistic models. NeurIPS 2020\n\nSong et al. Score-based generative modeling through stochastic differential equations. ICLR 2021.\n\n* **To include inverse problems in experiments for which DDPM has advantage**\n\nWe appreciate the reviewer's suggestion on the application to inverse problems. For the application to image data, image restoration problems can be cast as a conditional generation task. In this work, we have investigated conditional generation in Section 5.5 exemplified by graph data, where we applied the proposed model to learn from data the conditional distribution $X|Y$, where $X$ and $Y$ are nodal descriptors and labels on a graph, respectively. We demonstrated that JKO-iFlow outperforms the previous baseline in terms of generative quality, NLL, and MMD metrics. This setting may also be viewed as an inverse problem. The extension to other inverse problems is an interesting future direction.\n\n* **Can probability trajectories be stiff and require small step size to ensure stability and accuracy of the numerical ODE?**\n\nWe thank the reviewer for asking the interesting question. First, in the context of JKO-iFlow, a neural ODE block is used to parametrize the $k$-th step transport map $T_k$, which pushwards from $p_{k-1}$ to $p_k$. Theoretically, as long as both distributions have smooth densities, there exists a smooth transport map in between [Villani 2021]. Meanwhile, since the JKO scheme in the case of a small step size approximates the solution of a diffusion process FPE, which smoothes the distributions as time increases, we expect the discrete-time densities $p_k$ to also become more and more smooth as $k$ increases. In the case that the initial-time density (the data distribution) is not regular, this will make the flow at the early time irregular, and we observed that this is the place where “stiffness” may more often present. As the reviewer pointed out, this can require a small step size of the neural ODE.\n\nThe time-reparametrization introduced in Section 4.2 can be viewed as an adaptive way to adjust the step size to overcome this difficulty, and we have shown that this improves the performance of the JKO-iFlow model in experiments, specifically on toy data and MINIBOONE tabular data (see Sections 5.2 and 5.3). There are other works that address ODE stiffness and improve the stability of neural ODE, such as [Kim et al., 2021], which can be readily incorporated into our framework if needed.\n\nRef.\n\nC´edric Villani. Topics in optimal transportation, volume 58. American Mathematical Soc., 2021.\n\nKim S, Ji W, Deng S, et al. Stiff neural ordinary differential equations. Chaos: An Interdisciplinary Journal of Nonlinear Science, 2021, 31(9).\n\n* **FPE also appears in mean-field games, and can JKO-ifow work well on these other problems?**\n\nWe thank the reviewer for suggesting the application of the JKO-iFlow model to broader problems of mean-field games (MFG). The current work focuses on the continuous normalizing flow (CNF) setting and the application to generative models. Under the framework of MFG from time $0$ to $T$, if one sets the terminal cost as the $KL$-divergence between the density $\\rho_T$ and $p_Z = N(0, I)$, then the CNF problem is a special case of the MFG problem. We think that the JKO-iFlow model potentially can be applied to other MFG problems, and such extensions would be a very interesting future direction."}, {"review_id": "nu5XDVGxbS", "reviewer": "Reviewer_xTiV", "summary": "This paper presents an innovative method for enhancing the stability of trajectories in CNF (Continuous Normalizing Flow) models. The authors propose incorporating the JKO scheme, which introduces regularization between the current density and the base distribution. The core concept revolves around leveraging the Wasserstein distance to learn an optimal transport map that guides the data distribution towards a normal equilibrium. To assess the effectiveness of the model, experiments are conducted on various scenarios, including toy examples, tabular data, and datasets characterized by higher dimensions like CIFAR or MNIST.", "questions": "I do not have any particular questions for the authors. ", "limitations": "The main limitation of the proposed approach is the lack of capabilities to generate good-quality of image data, but this group of methods is rather not designed for such problems. Moreover, the evaluation may be a bit broader and include some other applications mentioned in the \"weaknesses\" section. Besides that, I admire the contribution of the work. and I think it should be accepted to the conference. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The idea behind the paper is novel and non-trivial. \n- Optimizing the velocity field instead of the transport map using JKO-Scheme is intriguing.\n- The procedure of training the model given the Algorithm is very efficient compared to reference baselines.\n- The method that can be seen as a particular case of regularisation of CNF models achieves outstanding results on tabular benchmark datasets.\n- The authors consider higher-dimension cases by conducting experiments on CIFAR. It is not common for such papers. \n", "weaknesses": "- The proposed method still does not provide good-quality image samples. This group of models is rather dedicated to obtaining high values of NLL and modeling the distributions of low-dimensional data. \n- The variety of experiments conducted in this work may be richer. In particular, it would be interesting to see how the proposed approach works in some scenarios where CNFs perform well, like modeling multidimensional probabilistic regression, using the model in latent space of VAEs, or as a plugin model to large models (like StyleFlow). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Please refer to the global response for the question on image quality and additional results of image generation using VAE latent space. \n\n* **(Weakness 2) The variety of experiments conducted in this work may be richer. In particular, it would be interesting to see how the proposed approach works in some scenarios where CNFs perform well, like modeling multidimensional probabilistic regression, using the model in latent space of VAEs, or as a plugin model to large models (like StyleFlow).**\n\nWe thank the reviewer for suggesting a rich class of applications which are suitable scenarios to apply CNF models. For the latent space of VAEs, our real image generation experiments (CIFAR10 and the newly added Imagenet-32) are actually obtained by applying the JKO-iFlow model in the VAE latent space, see more in the global response above. We agree that other problems like multi-dimensional probabilistic regression [Chen et al., 2018] and plugins to deep architectures like StyleFlow [Abdal et al., 2021] are interesting extensions. We will add all this to the discussion of future directions.\n\nRef.\n\nChen R T Q, Rubanova Y, Bettencourt J, et al. Neural ordinary differential equations. NeurIPS 2018.\n\nAbdal R, Zhu P, Mitra N J, et al. Styleflow: Attribute-conditioned exploration of stylegan-generated images using conditional continuous normalizing flows. ACM Transactions on Graphics (ToG), 2021."}, {"review_id": "7i8w23DNUX", "reviewer": "Reviewer_3zMK", "summary": "The paper introduces a novel approach to train continuous normalizing flows/score-based diffusion models that exploits the Wasserstein gradient flow theory and the JKO iterative scheme.\n\nThe main idea is to approximate the density evolution of the variance-preserving forward dynamics of a diffusion model model (i.e. an Ornstein -Uhlenbeck process) with a series of deterministic transport maps. Under some conditions, the resulting deterministic dynamics is invertible and can then be used to map the stationary distribution back to the (implicitly learned) data distribution.\n\nIn practice, this is very similar to what is done in ordinary score-based diffusion models, but the networks are trained in a radically different way since with the JKO approach the score is trained indirectly through a deterministic algorithm.\n", "questions": "- Can you elaborate on the connections between the JKI iterative training and the score-matching training of diffusion models? It would be informative to discuss differences and similarities and give an overview of the strengths and weaknesses of each approach.\n\n- I would like to see the FID scores on MNIST and CIFAR10.\n", "limitations": "Nothing to discuss.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "I find this work exciting. The approach is very original as introduces a brand new way to train score-based models, which is backed by serious theoretical arguments. This both increases my confidence in the algorithm and opens the door for possibly very fruitful connection between several important research areas.\n\nThe paper is well-explained, although it relies on several complex concepts that could alienate some readers. However, I think that in this case the mathematical complexity is unavoidable and gives added value. \n\nI appreciate that the authors not not solely focus on the mathematical framework but also discuss several important issues in the algorithmic implementation.\n\nThe experimental section is convincing and rather comprehensive.  The results on the low-dimensional datasets are good when compared with other state-of-the-art methods. The results on images are less impressive but overall convincing. While, in general, the application to generative computer vision needs to be fully flashed out, I agree with the authors that this goes beyond the scope of this work.\n", "weaknesses": "I think that, generally speaking, the paper does not have major weaknesses. However, the part of the experiment section dedicated to image generation would have benefited from having a larger scope, for example by including experiments on CelebA HD and ImageNet. \n\nThe selection of evaluation metrics should be complemented with FID scores, since this is the standard metric in the image generation literature and it is generally more reliable than the BPD score. Note that I do not think that having a lower FID score than the baseline is an argument for rejection, given the theoretical and algorithmic novelty of the work. \n\nTher exposition would benefit from a more extended section connecting the JKO approach to the standard score-matching diffusion theory. This would greatly help readers coming from the diffusion literature, who in my opinion are one of the most important audiences for this work. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Please refer to the global response for the questions on experiments of image generation and FID scores. \n\n* **(Weakness 3rd paragraph) The exposition would benefit from a more extended section connecting the JKO approach to the standard score-matching diffusion theory. This would greatly help readers coming from the diffusion literature, who in my opinion are one of the most important audiences for this work.**\n\nWe thank the reviewer for the suggestion. The connection between the JKO approach and the standard score-matching diffusion model is summarized below, and we will add it to the revised manuscript. \n\nFirst, it is known that the diffusion SDE for the stochastic process formed by the particles is connected to a deterministic ODE via the FPE (Fokker-Planck Equation). Specifically, the solution of the FPE Eqn (3) of the (variance preserving) SDE is equivalent to that of the continuity equation (CE) of the ODE Eqn (2) when the velocity field $f(x,t)$ is set to Eqn (9). While this relation has been used in the reverse-time (generating) process of score-matching diffusion model [Song et al., 2021], the JKO approach here is trying to directly learn a discrete-time ODE forward process and avoids sampling in training. The forward process approximates the CE and equivalently the FPE when step-size is small thanks to the JKO theory [Jordan et al., 1998]. \n\nMeanwhile, unlike the score-based diffusion model, which tries to learn the score $\\nabla \\log \\rho_t$, or equivalently the velocity field $f(x,t)$ that leads to exact solution of the FPE, with finite step size the JKO solution only approximates the solution of FPE. While such approximate sequence of densities $p_k$ do not have interpretation in physics, for the problem of normalizing flow it actually can provide a good CNF model without learning the exact solution of FPE: this is because as long as the last $p_k$ can be close enough to the normal density, and the transport in each learned residual block $T_k$ can be accurately inverted, this will guarantee the accuracy of the generating the data density in the reverse process. In this view, the JKO approach can explore the CNF models that transport from normal density to data density and back without requiring the intermediate process to exactly follow the diffusion process. The block-wise training scheme also allows savings and flexibility both in neural network architecture and in the training algorithm. \n\nRef:\n\nSong et al. Score-Based Generative Modeling through Stochastic Differential Equations. ICLR. 2021.\n\nJordan, Richard, David, Kinderlehrer, and Felix, Otto. \"The Variational Formulation of the Fokker–Planck Equation\". SIAM Journal on Mathematical Analysis 29, no.1 (1998): 1-17.\n\n* **(Question 1) Can you elaborate on the connections between the JKO iterative training and the score-matching training of diffusion models? It would be informative to discuss differences and similarities and give an overview of the strengths and weaknesses of each approach.**\n\nThanks for the suggestion. We will explain the difference and similarities of the two in the “connection” section as in the answer to the previous question, and also add the following summary of strengths and weaknesses to the draft.\n\nThe score-matching training has the advantage of more efficient optimization of the neural network that parametrizes the score function, thanks to the score-matching objective. In comparison, the current neural ODE model used in the JKO approach involves more expensive backpropagation of the neural network due to the ODE solver and the Hutchinson projection which undermines the scalability to higher dimensionality. Yet by paying this price, the CNF model has the advantage of optimizing the likelihood as the training objective and the ability to compute the likelihood by time integration. \n\nHowever, in training and generation the JKO approach has the advantage that it computes deterministic dynamics in both forward and reverse processes, avoiding simulating the SDE trajectories (injecting noise). The ODE integrator in JKO approach allows us to choose a much larger step size in the discrete time process compared to what is needed by the SDE simulation of score-based models. \n\nAnother advantage of the JKO approach lies in the block-wise training. In addition to the savings of computation and memory, the block-wise training of the JKO-iflow model has the potential flexibility that different blocks may have different architectures, which may be advantageous in certain settings. This type of flexibility may not be easy to implement in the end-to-end training of most score matching diffusion models. Finally, the JKO scheme is general and may be combined with other flow-based backbone models to overcome the computational issue, possibly with specific designs for certain applications, see more in the answer to R1.\n"}]

Author rebuttal:
Thanks for the constructive comments and feedback provided by all the reviewers. The common questions are about the experiments on image generation, which we first address here. The additional questions and comments of each reviewer are addressed in the specific responses below. R1 = Reviewer NZGt, R2 = Reviewer uXKQ, R3 = Reviewer xTiV, R4 = Reviewer 3zMK.

* **Quality of generated images, quantitative evaluation (FID), and additional image examples. [R1 Question 1, R3 Weakness 1 & Limitations, R4 Question 2]**

We have conducted additional experiments on CIFAR10. Please see the attached PDF for generated images by the final model, which show improved image quality. The model achieves an FID of 29.10. The improvement from the initial submission comes from using a larger VAE model following [Rombach et al., 2022] and longer training time (24 hours on one A100 GPU, the VAE training is extra). We have also computed the FID for MNIST, which is 7.95 using JKO-iFlow in the code space of a pre-trained VAE.

To enlarge the image generation experiments, we further examined our model on ImageNet-32, which consists of 1.28 million training images. The generated images are also shown in the attached PDF, and the model achieves an FID of 20.10 after training for 30 hours on one A100 GPU. We are aware that these FIDs do not achieve the SOTA performance of some diffusion-based models, but these results are obtained with less computation and this image generation performance is better than most CNF baselines. We will incorporate these results in the paper revision.

* **Extension to larger images [R1 Weakness 1, R2 Weakness (1st-2nd sentences), R4 Weakness paragraphs 1 & 2]**

We thank the reviewers for the suggestion, and larger images are certainly an important application direction of the proposed approach. Our current method for image generation is by applying a flow model in the latent VAE space. Note that generative models in the latent space, like StableDiffusion [Rombach et al., 2022], obtain SOTA image generation and are popular approaches for many industry applications. Thus, the VAE+flow approach has the potential to generate larger images. 

Meanwhile, we do think the extension to larger images would be more efficient if some fundamental development of the neural ODE can be done in the first place. A main reason is that the neural ODE backbone model currently taken by this work does not have a strong computational scalability to high-dimensional space due to the expensive backpropagation via ODE solver and the Hutchinson projection to compute $div(f)$ in the likelihood-based training. The current JKO approach is for general data and not specifically designed for large image generation tasks (as has been pointed out by R3). To extend the JKO approach to larger images, one can proceed by developing a more efficient neural ODE to scale to higher dimensional latent or pixel space or by combining the JKO scheme with other backbone models more suitable for image generation tasks, see more in the answer to R1 below. We leave these developments to future work and will add all this to the discussion.

Ref.

Rombach R, Blattmann A, Lorenz D, et al. High-resolution image synthesis with latent diffusion model. CVPR. 2022

Meta-review: In this paper, the authors introduce a neural ODE based formulation of a normalizing flow generative model.  The architecture of the formulation is composed of residual blocks that correspond to JKO Steps from a particular transport problem.  The authors provide a block-wise procedure to train this model and justify that it is more computationally efficient than end-to-end approaches in the literature.  The resulting approach is shown to have competitive or better performance compared to existing flow and diffusion models on image datasets that are small in image size.  The papers strengths include the signifiant novelty and motivation of a JKO inspired normalizing flow formulation, the computational efficiency of the proposed approach, and its competitive performance relative to related methods.  Reviewers noted that the method was tested on relatively small images, and the paper could be improved by evaluation of images of a larger size.  Given the difficulty of conducting such experiments on larger images, the strengths of the paper outweigh this drawback.
Decision: Accept (spotlight)