Title: Normalizing flow neural networks by JKO scheme
Venue: NeurIPS 2023
Accepted: yes
Keywords: Normalizing flow

invertible neural networks

JKO scheme
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

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 JKOiFlow, 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.

1 Introduction

Generative models have wide applications in statistics and machine learning to infer data-generating distributions and to sample from the model distributions learned from the data. In addition to widely used deep generative models such as variational auto-encoders (V AE) Kingma and Welling [2014, 2019] and generative adversarial networks (GAN) Goodfellow et al. [2014], Gulrajani et al. [2017], normalizing flow models Kobyzev et al. [2020] have been popular and with a great potential. The flow model learns the data distribution via an invertible mapping F between the data density p X in R d and the target standard multivariate Gaussian density p Z , Z ∼ N (0 , I d ) . While flow models, once trained, can be utilized for efficient data sampling and explicit likelihood evaluation, training of such models is often difficult in practice. To alleviate such difficulties, many prior works [Dinh et al., 2015, 2017, Kingma and Dhariwal, 2018, Behrmann et al., 2019, Ruthotto et al., 2020, Onken et al., 2021], among others, have explored

(a) JKO-iFlow

(b) usual CNF

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

Yao Xie

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

In this work, we propose JKO-iFlow, an invertible normalizing flow network that unfolds the Wasserstein gradient flow via a neural ODE model inspired by the JKO-scheme Jordan et al. [1998]. The JKO scheme, see (5), can be viewed as a proximal step to minimize the Kullback-Leibler (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 a smooth trajectory of density evolution, see Figure 1(a). 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 a 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 reparameterization with progressive refinement in computing the flow network, where each block corresponds to a representative point along the density evolution trajectory in the space of probability measures. 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., 2020, Song 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 Johnson 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 p Z having a parametrized potential V , exemplified by the application to the conditional generation setting where p Z 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 a 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.

· We show that JKO-iFlow greatly reduces memory and computational cost when achieving competitive 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)[Kingma and Welling, 2014, 2019]. Apart from known training difficulties (e.g., mode collapse [Salimans et al., 2016] and posterior collapse [Lucas et al., 2019]), these models do not provide likelihood or inference of data density. The normalizing flow framework [Kobyzev et al., 2020] has been extensively developed, including continuous flow [Grathwohl et al., 2019], Monge-Ampere flow [Zhang et al., 2018], discrete flow [Chen et al., 2019], and extension to non-Euclidean data [Liu et al., 2019, Mathieu and Nickel, 2020, Xu et al., 2022]. Efforts have been made to develop novel invertible mapping structures [Dinh et al., 2017, Papamakarios et al., 2017] and regularize the flow trajectories by transport cost [Finlay et al., 2020, Onken et al., 2021, 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., 2021] 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 to compute the likelihood; using the same idea of backward SDE, Zhang and Chen [2021] proposed joint training of forward and backward neural SDEs. Theoretically, latent diffusion Tzen and Raginsky [2019b,a] was used to analyze neural SDE models. The current work focuses on a neural ODE model where the deterministic vector field f ( x, t ) is to 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, Song et al. [2021] derived neural ODE model from the learned score function of the diffused data marginal distributions for all t . We experimentally obtain competitive or improved performance against the diffusion model on simulated two-dimensional and high-dimensional tabular data.

JKO-inspired deep models have been studied in several recent works. [Bunne et al., 2022] reformulated the JKO step for minimizing an energy function over convex functions. JKO scheme has also been used to discretize Wasserstein gradient flow to learn a deep generative model in [Alvarez-Melis et al., 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 et 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 a residual network but identified difficulty in calculating the push-forward distribution. The approach in [Fan et al., 2022] also relies on a variational formulation, which requires training an additional network similar to the discriminator in GAN using inner loops. The idea of progressive additive learning in training generative ResNet, namely training ResNet block-wisely by a variational loss, dates back to Johnson and Zhang [2019] under the GAN framework. Our method trains an invertible neural-ODE flow network that flows from data density to the normal one and backward, which enables the computation of model likelihood as in other neural-ODE approaches. The objective in each JKO step to minimize KL divergence can also be computed directly without any inner-loop training, see Section 4.

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, 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 Vanden-Eijnden,

2023] matched the velocity field from interpolated distributions between initial and terminal ones; [Boffi and Vanden-Eijnden, 2023] proposed to learn the score ∇ log ρ t ( x ) (where ρ t solves the FPE and is unknown a priori ) to solve high-dimensional FPE. Using Stein's identity (which is equivalent to the derivation in Section 3.2), the step-wise training objective in [Boffi and Vanden-Eijnden, 2023] 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 et al., 2020] and studied via a self-consistency equation in [Shen et al., 2022]. While our approach learns an 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., 2023] proposed a 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 KL-divergence 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. [2017] established approximation by composition of Barron functions [Barron, 1993]; Bailey and Telgarsky [2018] developed space-filling approach, which was generalized in Perekrestenko et al. [2020, 2021]; Lu and Lu [2020] constructed a deep ReLU network with guaranteed approximation under integral probability metrics, using techniques of empirical measures and optimal transport. These results show that deep neural networks can provably transport one source distribution to a target one with sufficient model capacity under certain regularity conditions of the pair of densities. 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 ρ ( x, t ) such that ρ ( x, 0) = p X and as t increases ρ ( x, t ) approaches p Z ∼ N (0 , I d ) [Tabak and Vanden-Eijnden, 2010]. Given an initial distribution ρ ( x, 0) , such a flow typically is not unique. We consider when the flow is induced by an ODE of x ( t ) in R d

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

Ornstein-Uhlenbeck (OU) process . Consider a Langevin dynamic denoted by the SDE dX t = -∇ V ( X t ) dt + √ 2 dW t , where V is the potential of the equilibrium density p Z . We focus on the case of normal equilibrium, that is, V ( x ) = | x | 2 / 2 and then p Z ∝ e -V . In this case, the process is known as the (multivariate) OU process. Suppose X 0 ∼ p X , and let the density of X t be ρ ( x, t ) also denoted as ρ t ( · ) . The Fokker-Planck equation (FPE) describes the evolution of ρ t towards the equilibrium p Z as follows, where V ( x ) := | x | 2 / 2 ,

Under generic conditions, ρ t converges to p Z exponentially fast. For Wasserstein-2 distance and the standard normal p Z , classical argument gives that (take C = 1 in Eqn (6) of Bolley et al. [2012])

JKO scheme . The seminal work Jordan et al. [1998] established a time discretization scheme of the solution to (3) by the gradient flow to minimize KL( ρ || p Z ) under the Wasserstein-2 metric in probability space. Denote by P the space of all probability densities on R d with a finite second moment. The JKO scheme computes a sequence of distributions p k , k = 0 , 1 , · · · , starting from p 0 = ρ 0 ∈ P . With step size h > 0 , the scheme at the k -th step is written as

where F [ ρ ] := KL( ρ || p Z ) . It was proved in Jordan et al. [1998] that as h → 0 , the solution p k converges to the solution ρ ( · , kh ) of (3) for all k , and the convergence ρ ( h ) ( · , t ) → ρ ( · , t ) is strongly in L 1 ( R d , (0 , T )) for finite T where ρ ( h ) is piece-wise constant interpolated from p k .

3 JKO scheme by neural ODE

Given i.i.d. observed data samples X i ∈ R d , i = 1 , . . . , N , drawn from some unknown density p X , the goal is to train an invertible neural network to transport the density p X to an a priori specified density p Z in R d , where each data sample X i is mapped to a code Z i . A prototypical choice of p Z is the standard multivariate Gaussian N (0 , I d ) . In this work, we leave the potential of p Z abstract and denote by V , that is, p Z ∝ e -V and V ( x ) = | x | 2 / 2 for normal p Z . By a slight abuse of notation, we denote by p X and p Z both the distributions and the density functions of data X and code Z respectively.

3.1 Objective of JKO step

We are to specify f ( x, t ) in the ODE (1), to be parametrized and learned by a neural ODE, such that the induced density evolution of p ( x, t ) converges to p Z as t increases. We start by dividing the time horizon [0 , T ] into finite subintervals with step size h , let t k = kh and I k +1 := [ t k , t k +1 ) . Define p k ( x ) := p ( x, kh ) , namely the density of x ( t ) at t = kh . The solution of (1) determined by the vector-field f ( x, t ) on t ∈ I k +1 (assuming the ODE is well-posed [Sideris, 2013]) gives a one-to-one mapping T k +1 on R d , s.t. T k +1 ( x ( t k )) = x ( t k +1 ) and T k +1 transports p k into p k +1 , i.e., ( T k ) # p k -1 = p k , where we denote by T # p the push-forward of distribution p by T , such that ( T # p )( A ) = p ( T -1 ( A )) for a measurable set A . In other words, the mapping T k +1 is the solution map of the ODE from time t k to t k +1 .

Suppose we can find f ( · , t ) on I k +1 such that the corresponding T k +1 solves the JKO scheme (5), then with small h , p k approximates the solution to the Fokker-Planck equation 3, which then flows towards p Z . By the Monge formulation of the Wasserstein-2 distance between p and q as W 2 2 ( p, q ) = min T : T # p = q E x ∼ p ∥ x -T ( x ) ∥ 2 , solving for the transported density p k by (5) is equivalent to solving for the transport T k +1 by

where F [ T ] = KL( T # p k || p Z ) . The equivalence between (5) and (6) is proved in Lemma A.1.

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

Proposition 3.1. Given p k , up to a constant c independent from f ( x, t ) on t ∈ I k +1 ,

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

where x ( t k +1 ) = x ( t k ) + ∫ t k +1 t k f ( x ( s ) , s ) ds . Taking a neural ODE approach, we parametrize { f ( x, t ) } t ∈ I k +1 as a residual block with parameter θ k +1 , and then (8) is reduced to minimizing over θ k +1 . This leads to a block-wise learning algorithm to be introduced in Section 4, where we further allow the step-size h to vary for different k as well.

3.2 Infinitesimal optimal f ( x, t )

In each JKO step of (8), let p = p k denote the current density, q = p Z be the target equilibrium density. In this subsection, we show that the optimal f in (8) with small h reveals the difference between score functions between target and current densities. Thus, minimizing the objective (8) searches for a neural network parametrization of the score function ∇ log ρ t implicitly, in contrast to diffusion-based models which learn the score function explicitly [Ho et al., 2020, Song et al.,

2021], e.g., via denoising score matching. At infinitesimal h , this is equivalent to solving the FPE by learning a deterministic transport equation as in [Boffi and Vanden-Eijnden, 2023, Shen et al., 2022].

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

where x ( h ) = x (0)+ ∫ h 0 f ( x ( s ) , s ) ds . In the limit of h → 0+ , formally, x ( h ) -x (0) = h f ( x (0) , 0)+ O ( h 2 ) , and suppose V of q is C 2 , V ( x ( h )) = V ( x (0)) + h ∇ V ( x (0)) · f ( x (0) , 0) + O ( h 2 ) . For any differentiable density ρ , the (Stein) score function is defined as s ρ = ∇ log ρ , and we have ∇ V = -s q . Taking the formal expansion of orders of h , the objective in (9) is written as

Note that E x ∼ p V ( x ) is independent of f ( x, t ) , and the O ( h ) order term in (10) is over f ( x, 0) only, thus the minimization of the leading term is equivalent to

where T q is known as the Stein operator [Stein, 1972]. The T q f in (11) echoes that the derivative of KL divergence with respect to transport map gives Stein operator [Liu and Wang, 2016]. The Wasserstein-2 regularization gives an L 2 regularization in (11). Let L 2 ( p ) be the L 2 space on ( R d , p ( x ) dx ) , and for vector field v on R d , v ∈ L 2 ( p ) if ∫ | v ( x ) | 2 p ( x ) dx < ∞ . One can verify that, when both s p and s q are in L 2 ( p ) , the minimizer of (11) is

This shows that the infinitesimal optimal f ( x, t ) equals the difference between the score functions of the equilibrium and the current density.

3.3 Invertibility of flow model and expressiveness

At time t the current density of x ( t ) is ρ t , the analysis in Section 3.2 implies that the optimal vector field f ( x, t ) has the expression as

With this f ( x, t ) , the Liouville equation (2) coincides with the FPE (3). This is consistent with the JKO scheme with a small h recovering the solution to the FPE. Under proper regularity condition of V and the initial density ρ 0 , the r.h.s. of (12) is also regular over space and time. This leads to two consequences, in approximation and in learning: Approximation-wise, the regularity of f ( x, t ) allows to construct a k -th residual block in the flow network to approximate { f ( x, t ) } t ∈ I k 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.

For learning, when properly trained with sufficient data, the neural ODE vector field f ( x, t ; θ k ) will learn to approximate (12). This can be viewed as inferring the score function of ρ t , and also leads to the invertibility of the trained flow net in theory: Suppose the trained f ( x, t ; θ k ) is close enough to (12); it will also have bounded Lipschitz constant. Then the residual block is invertible as long as the step size h is sufficiently small, e.g. less than 1 /L where L is the Lipschitz bound of f ( x, t ; θ k ) . In practice, we typically use smaller h 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.1).

4 Training of JKO-iFlow net

The proposed JKO-iFlow model allows progressive learning of the residual blocks in the neural-ODE model in a block-wise manner (Section 4.1). We also introduce two techniques to improve the training of the trajectories in probability space (Section 4.2), 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 i = x i (0) } n i =1 , the expectation E x ( t ) ∼ p k in (8) is replaced by the sample average over { x i ( kh ) } n i =1 which can be computed from the previous k blocks. Note that for each given x ( t ) = x ( t k ) , both x ( t k +1 ) and the integral of ∇· f in (8) can be computed by a numerical neural ODE integrator. Following previous works, we use the Hutchinson trace estimator [Hutchinson, 1989, Grathwohl et al., 2019] to compute the quantity ∇· f in high dimensions, and we also propose a finite-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 f θ k with trainable parameters θ k .

This leads to a block-wise training of the normalizing flow network, as summarized in Algorithm 1. The sequence of time stamps t k is given by specifying the time steps h k := t k +1 -t k , which we allow to differ across k . The choice of the sequence h k is initialized by a geometric sequence starting from h 0 with maximum stepsize h max , see Appendix B.1. In the special case where the multiplying factor is one, the sequence of h k gives a constant step size. The adaptive choice of h k with refinement (by adding more blocks) will be introduced in Section 4.2. Regarding the termination criterion Ter(k) in line 2 of Algorithm 1, we monitor the ratio E x ∼ p k -1 ∥ x -T k ( x ) ∥ 2 / E x ∼ p k -1 ∥ T k ( x ) ∥ 2 and terminate when it is below some threshold ϵ , set as 0.01 in all experiments. In practice, when training the k -th block, both the averages of ∥ x -T k ( x ) ∥ 2 and ∥ T k ( x ) ∥ 2 are computed by empirically averaging over the training samples (in the last epoch) at no additional computational cost. Lastly, line 5 of training a 'free block' (i.e., the block without the W 2 regularization) is to flow the push-forward density p L closer to the target density p Z , where the former is obtained through the first L blocks.

The block-wise training significantly reduces the memory and computational load since only one block is trained when optimizing (8) regardless of flow depth. Therefore, one can use larger training batches and potentially more expensive numerical integrators within a certain memory budget for higher accuracy. We also empirically observe that training each block using standard back-propagation (rather than the adjoint method in neural ODE) gives a comparable result at a lower cost. To ensure the invertibility of the trained JKO-iFlow network, we further break the time interval [ t k -1 , t k ) into 3 or 5 subintervals to compute the neural ODE integration, e.g.,

Algorithm 1 Block-wise JKO-iFlow training

Require: Time stamps { t k } , training data, termination criterion Ter and tolerance level ϵ , maximal number of blocks L max .

1: Initialize k = 1 .

2: while Ter(k) > ϵ and k ≤ L max do

3: Optimize f θ k upon minimizing (8) with mini-batch sample approximation, given { f θ i } k -1 i =1 . Set k ← k +1 .

4: end while

5: L ← k . Optional: Optimize f θ L +1 without W 2 regularization.

by Runge-Kutta-4. We empirically verify small inversion errors on test samples.

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 p k , k = 1 , . . . , L , associated with the L residual blocks of the proposed normalizing flow network. The two techniques are illustrated in Figure 2. Further details and illustrations of the approach can be found in Appendix B.

Trajectory reparameterization . We empirically observe fast decay of the movements W 2 2 ( T # p k , p k ) when h k is set to be constant, that is, initial blocks transport the densities much further than the later ones. This is consistent with the exponential convergence of the Fokker-Planck flow, see (4), 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 h k , we reparameterize the values of t k through an adaptive procedure based on the W 2 distance at each block. The procedure uses an adaptive approach to encourage the W 2 movement in each block to be more even across the L blocks, where the retraining of the trajectory can be potentially warm-started by the previous trajectory in iteration.

Progressive refinement . The performance of CNF models is typically improved with a larger number of residual blocks, corresponding to a smaller step size. The smallness of stepsize h also ensures the invertibility of the flow model, in theory and also in practice. However, directly training the model with a small non-adaptive stepsize h may result in long computation time and convergence to the normal density q Z only after a large number of blocks, where the choice of h k is not as efficient as after adaptive reparameterization. We introduce a refinement approach

W

2

Block

1

2

3

p

0

=

p

X

p

1

p

2

p

3

p

Z

W

2

Block

1

3

4

2

5

6

p

0

=

p

X

p

Z

p

1

p

2

p

3

p

4

p

5

p

6

Probability Trajectory

Movement

W

2

f

6

f

5

f

4

f

3

f

2

f

1

f

1

f

2

f

3

to increase the number of blocks progressively, where each time interval [ t k -1 , t k ) splits into two halves, and the number of blocks doubles after the adaptive reparameterization of the trajectory converges at the coarse level. The new trajectory at the refined level is again trained with adaptive reparameterization, where the residual blocks can be warm-started from the coarse-level model to accelerate the convergence. The trajectory refinement allows going to a smaller step size h k , which benefits the accuracy of the JKO-iFlow model including the numerical accuracy in integrating each neural ODE block.

5 Experiment

In this section, we examine the proposed JKO-iFlow model on simulated and real datasets, including both unconditional and conditional generation tasks. Codes are available at https://github.com/ hamrel-cxu/JKO-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-ODE models). The next two CNF models are IResNet [Behrmann et al., 2019] and IGNN [Xu et al., 2022], which are discrete in time (ResNet models). The diffusion model baseline is the score-based neural SDE [Song et al., 2021], 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 log-likelihood (NLL) metric, and the kernel maximum mean discrepancy (MMD) [Gretton et al., 2012a] metric, including MMD-1, MMD-m, and MMD-c for constant, median distance, and custom bandwidth respectively. The test threshold τ 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 α = 0 . 05 test level). See details in Appendix C.1. The computational cost is measured by the number of mini-batch stochastic gradient descent steps and the training

(a) True data

JKO-iFlow

τ

: 2.79e-4

, MMD-c:

2.73e-4

NLL

2.64

(b) FFJORD

3.88e-4

2.95

(c) OT-Flow

1.42e-3

3.30

(d) IGNN

3.14e-3

3.35

(e) ScoreSDE

6.90e-4

3.2

(f) Fractal tree

τ

: 3.12e-4

, MMD-c:

2.17e-4

NLL

2.20

(g) Olympic rings

τ

: 3.16e-4

, MMD-c:

2.36e-4

NLL

1.66

(h) Checkerboard

τ

: 3.09e-4

, MMD-c:

2.70e-4

NLL

3.59

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 1e-5 level or below, see Table A.1.

5.2 Two-dimensional toy data

The results on four toy datasets are shown in Figure 3, where the metrics of NLL and MMD-c are reported. Plots (a)-(e) compare JKO-iFlow with the alternative models on the dataset Rose, where both NLL and MMD-c metrics show the better performance of JKO-iFlow. Visually, the generated samples ˆ 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 a satisfactory generative performance. The results of JKO-iFlow in Figure 3 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 1 assesses the performance of JKO-iFlow and baseline methods on 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 we 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.4.

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 are shown in Figure A.5, where the W 2 movements across the blocks become more even 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 Imagenet32 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 have been shown to obtain state-of-the-art image generation performance, e.g., in StableDiffusion [Rombach et al., 2022]. Here we train a flow model instead of a score-based diffusion model in the latent space, and more details can be found in Appendix C.2.

The generated images are shown in Figure 4. The quantitative evaluation is by the Fr´ echet 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-art 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 A100 GPU, our experiments took 24 hours on CIFAR10 and 30 hours on Imagenet-32. Meanwhile, the image generation by JKO-iFlow model here obtains visually more appealing images and achieves lower

/143349488

00

SWS

76979

9

08YO

74357972

FIDs compared to most CNF baselines [Grathwohl et al., 2019, Behrmann et al., 2019, Chen et al., 2019, Finlay et al., 2020].

5.5 Conditional generation

The problem aims to generate input samples X given a label Y from the conditional distribution X | Y 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 | Y = k to the class-specific component in the Gaussian mixture of H | Y = k , see Figure A.7b and Appendix C.2.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 A.8, 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 in neural ODE in addition to the dimension reduction techniques by VAE as been explored here. Another possibility is to combine the JKO-iFlow scheme with other backbone flow models that are more suitable for the specific tasks. Meanwhile, it would be interesting to 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], a plug-in to deep architectures such as StyleFlow [Abdal et al., 2021], and the application to Mean-field Games [Huang et al., 2023].

Theoretically, the expressiveness of the flow model to generate a regular data distribution can be analyzed based on Section 3.3. To sketch a 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 ε -close to normality in T = log(1 /ε ) time, see (4). Reversing the time of the ODE then leads to an approximation of the initial density ρ 0 = p X by flowing backward in time from T to zero. Further analysis under technical assumptions is left to future work. During the time that this paper was being published, the convergence analysis of the proposed model was studied in Cheng et al. [2023].

References

Rameen Abdal, Peihao Zhu, Niloy J Mitra, and Peter Wonka. Styleflow: Attribute-conditioned exploration of stylegan-generated images using conditional continuous normalizing flows. ACM Transactions on Graphics (ToG) , 40(3):1-21, 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. Transactions on Machine Learning Research , 2022. ISSN 2835-8856.

Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning , pages 146-155. PMLR, 2017.

Miguel A Arcones and Evarist Gine. On the bootstrap of u and v statistics. The Annals of Statistics , pages 655-674, 1992.

Bolton Bailey and Matus J Telgarsky. Size-noise tradeoffs in generative networks. Advances in Neural Information Processing Systems , 31, 2018.

Andrew R Barron. Universal approximation bounds for superpositions of a sigmoidal function. IEEE Transactions on Information theory , 39(3):930-945, 1993.

Jens Behrmann, Will Grathwohl, Ricky TQ Chen, David Duvenaud, and J¨ orn-Henrik Jacobsen. Invertible residual networks. In International Conference on Machine Learning , pages 573-582. PMLR, 2019.

Adam Block, Youssef Mroueh, and Alexander Rakhlin. Generative modeling with denoising autoencoders and langevin sampling. arXiv preprint arXiv:2002.00107 , 2020.

Nicholas M Boffi and Eric Vanden-Eijnden. Probability flow solution of the fokker-planck equation. Machine Learning: Science and Technology , 4(3):035012, 2023.

Franc ¸ois Bolley, Ivan Gentil, and Arnaud Guillin. Convergence to equilibrium in wasserstein distance for fokker-planck equations. Journal of Functional Analysis , 263(8):2430-2457, 2012.

Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics , pages 6511-6528. PMLR, 2022.

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¨ orn-Henrik 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.17582 , 2023.

Micha¨ el Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. Advances in neural information processing systems , 29, 2016.

Pierre Degond and S Mas-Gallic. The weighted particle method for convection-diffusion equations. i. the case of an isotropic viscosity. Mathematics of computation , 53(188):485-507, 1989.

Pierre Degond and Francisco-Jos´ e Mustieles. A deterministic approximation of diffusion equations using particles. SIAM Journal on Scientific and Statistical Computing , 11(2):293-310, 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 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 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-resolution 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 Machine Learning , pages 6185-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¨ orn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ode: the world of jacobian and kinetic regularization. In International conference on machine learning , pages 3154-3164. PMLR, 2020.

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¨ olkopf, 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 large-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, Bernhard 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-Wei Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective on diffusion-based generative models and score matching. Advances in Neural Information Processing Systems , 34: 22863-22876, 2021.

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 Statistics - Simulation and Computation , 18:1059-1076, 1989.

Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A. Efros. Image-to-image translation with conditional adversarial networks. 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 5967-5976, 2017.

Rie Johnson and Tong Zhang. A framework of composite functional gradient methods for generative adversarial models. IEEE transactions on pattern analysis and machine intelligence , 43(1):17-32, 2019.

Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker-planck equation. SIAM journal on mathematical analysis , 29(1):1-17, 1998.

Diederik P. Kingma and Jimmy Ba. Adam: A 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 P. Kingma and Max Welling. An introduction to 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 on Learning Representations , 2021.

Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical Report 0, 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 Theory , pages 1271-1296. 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 on Learning Representations , 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 in neural information processing systems , 29, 2016.

Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Robert St¨ oter. Sliced-wasserstein 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 in neural information processing systems , 33: 3094-3105, 2020.

James Lucas, G. Tucker, Roger B. Grosse, and Mohammad Norouzi. Understanding posterior collapse in generative latent variable models. In DGS@ICLR , 2019.

Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions of fokkerplanck equations through gradient-log-density 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 , 34:15243-15256, 2021.

Derek Onken, S Wu Fung, Xingjian Li, and Lars Ruthotto. Ot-flow: Fast and accurate continuous normalizing flows via optimal transport. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, 2021.

George Papamakarios, Theo Pavlakou, and Iain Murray. Masked autoregressive flow for density estimation. Advances in 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 DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alch´ e-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32 , pages 8024-8035. Curran Associates, Inc., 2019.

Dmytro Perekrestenko, Stephan M¨ uller, and Helmut B¨ olcskei. Constructive universal highdimensional distribution generation through deep relu networks. In International Conference on Machine Learning , pages 7610-7619. PMLR, 2020.

Dmytro Perekrestenko, L´ eandre Eberhard, and Helmut B¨ olcskei. 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¨ orn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10684-10695, 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 , 117(17):9183-9193, 2020.

Tim Salimans, Ian Goodfellow, Wojciech 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, Ilmun Kim, M ˜ A©lisande Albert, B ˜ A©atrice Laurent, Benjamin Guedj, and Arthur Gretton. Mmd aggregated two-sample test. Journal of Machine Learning Research , 24(194):1-81, 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 C 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.

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 P 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 J. 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 G Tabak and Eric Vanden-Eijnden. Density estimation by dual ascent of the log-likelihood. 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 arXiv:1905.09883 , 2019a.

Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory , pages 3084-3114. PMLR, 2019b.

Alexander Vidal, Samy Wu Fung, Luis Tenorio, Stanley Osher, and Levon Nurbekyan. Taming hyperparameter tuning in continuous normalizing flows using the jko scheme. Scientific 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):454-467, 2022.

Dmitry Yarotsky. Error bounds for approximations with deep relu networks. Neural Networks , 94: 103-114, 2017.

Linfeng Zhang, Lei Wang, et al. Monge-Amp´ ere flow for generative modeling. arXiv preprint arXiv:1809.10188 , 2018.

Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in Neural Information Processing Systems , 34:16280-16291, 2021.

A Proofs

A.1 Proofs in Section

Lemma A.1. Suppose p and q are two densities on R d in P , the following two problems

have the same minimum and

(a) If T ∗ : R d → R d is a minimizer of (14) , then ρ ∗ = ( T ∗ ) # p is a minimizer of (13) .

(b) If ρ ∗ is a minimizer of (13) , then the optimal transport from p to ρ ∗ minimizes (14) .

Proof of Lemma A.1. Let the minimum of (14) be L ∗ T , and that of (13) be L ∗ ρ .

Proof of (a): Suppose L T achieves minimum at T ∗ , then T ∗ is the optimal transport from p to ρ ∗ = ( T ∗ ) # p because otherwise L T can be further improved. By definition of L ρ , we have L ∗ T = L T [ T ∗ ] = L ρ [ ρ ∗ ] ≥ L ∗ ρ . We claim that L ∗ T = L ∗ ρ . Otherwise, there is another ρ ′ such that L ρ [ ρ ′ ] < L ∗ T . Let T ′ be the optimal transport from p to ρ ′ , and then L T [ T ′ ] = L ρ [ ρ ′ ] < L ∗ T , contradicting with that L ∗ T is the minimum of L T . This also shows that L ρ [ ρ ∗ ] = L ∗ T = L ∗ ρ , that is, ρ ∗ is a minimizer of L ρ .

Proof of (b): Suppose L ρ achieves minimum at ρ ∗ . Let T ∗ be the OT from p to ρ ∗ , then E x ∼ p | x -T ∗ ( x ) | 2 = W 2 ( p, ρ ∗ ) 2 , and then L T [ T ∗ ] = L ρ [ ρ ∗ ] = L ∗ ρ which equals L ∗ T as proved in (a). This shows that T ∗ is a minimizer of L T .

Proof of Proposition 3.1, Given p k being the density of x ( t ) at t = kh , recall that T is the solution map from x ( t ) to x ( t + h ) . We denote ρ t := p k , and ρ t + h := T # p k . By definition,

Because p Z ∝ e -V , V ( x ) = | x | 2 / 2 , we have log p Z ( x ) = -V ( x ) + c 1 for some constant c 1 . Thus

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

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

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

Inserting back to (17), we have

The first term is determined by ρ t = p k , and thus is a constant c 2 independent from f ( x, t ) on t ∈ [ kh, ( k +1) h ] . Together with (16), we have shown that

which proves (7).

Derivation of (18): by chain rule,

B Technical details of Section 4

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

B.1 Initial choice of t k in Algorithm 1

We first initialize the choice by the following scheme

where the base stepsize h 0 , the multiplying factor ρ ≥ 1 , and the maximum stepsize h max ≥ h 0 are three hyper-parameters to be specified by the user. When ρ = 1 , the sequence of h k in (19) corresponds to constant stepsize h k ≡ h 0 .

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

B.2 Hutchinson trace via finite difference

We propose a finite-difference estimator of ∇· f ( x, t ) = Tr ( J f ( x )) , where J f ( x )) is the Jacobian of f at x . This is based on the Hutchinson trace estimator [Hutchinson, 1989], which states that Tr ( J f ( x )) = E p ( ϵ ) [ ϵ T J f ( x ) ϵ ] . Here, p ( ϵ ) is a distribution in R d satisfying E [ ϵ ] = 0 and Cov ( ϵ ) = I (e.g., a standard Gaussian). Existing CNF works estimate the expectation by sampling ϵ ∼ p ( ϵ ) and computing the vector-Jacobian product J f ( x ) ϵ using reverse-mode automatic differentiation. The product J f ( x ) ϵ can be computed for approximately the same cost as evaluating f [Grathwohl et al., 2019].

Now, given a fixed ϵ , we have J f ( x ) ϵ = lim σ → 0 f ( x + σϵ ) -f ( x ) σ , which is the directional derivative of f along the direction ϵ . Thus, given a fixed small σ > 0 , we propose the following finite-difference estimator of ∇· f ( x, t ) :

The finite-difference approximation of J f ( x ) ϵ requires only one additional function evaluation of f at x + σϵ and avoids the computation of automatic differentiation. We empirically found that training with (20) on high-dimensional examples (e.g., d > 100 as in image examples) can be approximately 1.5 ∼ 2 times faster than existing CNF approaches relying on reverse-mode automatic differentiation. In all experiments, we let σ = σ 0 / √ d with σ 0 = 0 . 02 .

B.3 Trajectory improvement illustrated in vector space

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

Let the space be R d , and F : R d → R be a differential landscape. Given a sequence of positive stepsize h k , one can compute a sequence of representative points x k which discretizes the gradient descent trajectory of minimizing F (towards a local minimum x ∗ ). Specifically, we have

starting from some fixed x 0 .

B.3.1 Trajectory reparameterization

Starting from an initial sequence of step-size { h k } L k =1 as in Section B.1, the reparameterization technique implements an iterative scheme to update the sequence of h k adaptively. We call the j -th iteration 'Iterj ', and denote the sequence by { h ( j ) k } L k =1 . The corresponding solution points x k by solving (21) with h ( j ) k are denoted as x ( j ) k . In addition to the constant h max , we need another algorithmic parameter η ∈ (0 , 1) , and we set η = 0 . 3 for the vector-space example in computation.

We always have x ( j ) 0 = x 0 for all 'Iterj '. In the j -th iteration of the parametrization, we compute the following

Compute the arclength as

Compute the average arclength

Update the stepsize for k = 1 , . . . , L as

Solve (21) using h ( j +1) k to obtain x ( j +1) k .

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

B.3.2 Progressive refinement

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

Given a sequence of { h k } L k =1 and the representative points { x k } L k =1 , the refinement scheme computes a refined trajectory having k = 1 , · · · , 2 L :

Compute the refined stepsizes as

Compute the representative points ˜ x k by solving (21) using ˜ h k , possibly warm-starting by ˜ x 2 k = x k and ˜ x 2 k -1 = ( x k + x k -1 ) / 2 .

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

B.4 Trajectory improvement in probability space

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

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

where the transport mapping T k ( x ) = x + ∫ t k t k -1 f θ k ( x ( s ) , s ) ds and can be computed from the k -th block. In practice, the expectation E x ∼ p k -1 in (23) is computed by a finite-sample average on the training set.

At last, the optimal warm-start of p k in the refinement is implemented by inheriting the parameters θ k of the trained blocks.

C Experimental details

C.1 Quantitative evaluation metrics

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

arclengths mean& std

initial:0.344 & 0.097

Iter-8:0.357&0.018

r-lter-3:0.184 & 0.009

Iter-1:0.347 & 0.078

Iter-9:0.357 &0.015

r-lter-4:0.184 & 0.007

Iter-2:0.350&0.062

Iter-10:0.357& 0.013

r-lter-5:0.184 &0.006

Iter-3:0.352&0.051

Iter-11:0.358 & 0.010

r-lter-6:0.184 & 0.004

Iter-4:0.353 &0.041

Iter-12:0.358&0.008

r-lter-7:0.184 & 0.003

Iter-5:0.355&0.034

refine:0.183&0.023

r-lter-8:0.184&0.003

Iter-6:0.355&0.027

r-lter-1:0.183 & 0.017

r-lter-9:0.184& 0.002

Iter-7:0.356 & 0.022

r-lter-2:0.184 & 0.012

r-lter-10:0.184 & 0.002

0.4

0.2

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

0.45

time ∑hi

Computed trajectorywithL=16

3.0

2.5

2.0

1.5

1.0

0.5

0.0

X16:[-1.94 0.244]

-0.5

True minumum: [-1.911, 0.103]

-1.0

5.0

4.5

4.0

3.5

-3.0

2.5

2.0

1.5

C.1.1 NLL metric

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

Specifically, our model flows from t 0 = 0 to t L +1 = T , where x (0) ∼ p 0 the data distribution, and we train the flow model to make x ( T ) follow a normal density q . The model density at data sample x is expressed as ρ ( x, 0) = ( T -1 θ ) # q ( x ) where T -1 θ is the inverse model flow mapping from x ( T ) to x (0) . Thus the model log-likelihood can be expressed as

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

where

starting from x (0) = x . Both the integration of f θ k and ∇· f θ k are computed using the numerical scheme of neural ODE. We report NLL in the natural unit of information (i.e., log with base e , known as 'nats') in all our experiments.

C.1.2 MMDmetrics

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 := { x i } N i =1 and ˜ X := { ˜ x j } j M =1 and a kernel function k ( x, ˜ x ) , the (squared) kernel MMD [Gretton et al., 2012a] is defined as

When a generative model is trained, we 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], Sutherland et al. [2017], Schrab et al. [2023]. In all experiments here, we use the Gaussian kernel k ( x, ˜ x ) = exp {-∥ x -˜ x ∥ 2 / 2 h 2 } to stay consistent with reported baselines from [Onken et al., 2021], where h > 0 is the bandwidth parameter. We use three ways of setting the bandwidth parameter h :

Constant bandwidth: h = h c = 1 . The resulting MMD is denoted as 'MMD-1'.

Median bandwidth [Gretton et al., 2012a]: let h = h m be the median of ∥ x i -x j ∥ for all distinct i and j . The median distance is computed from the dataset X . The resulting MMD is denoted as 'MMD-m'.

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

On all datasets, we use at most N = 10 K test samples as X , and for each trained model, we generate M = 10 K 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 is significant when above a test threshold τ , which is defined as the upper ( 1 -α )-quantile of the distribution of the MMD

0.4

0.2

0.0

3

6

9

12

15

18

21

24

Initial

Iter 1

Iter 2

Iter3

(a) Per-block

W

2

2

over reparameterization iterations.

MMDThresholdt:3.10e-04

MMD:3.01e-4,NLL:2.65

MMD:3.09e-4,NLL:2.64

(b) Results at initial training (middle) and Iter 3 (right). MMD and NLL values are shown in the title.

W²(pk-1.Pk)

1.0

0.5

0.0

2

4

5

8

10

12

14

16

Initial

Iter 2

Iter 4

refine

r-lter1

Iter 1

Iter 3

(a) Per-block W 2 2 over reparameterization iterations and refinement ('r-Iter 1' means one reparameterization iteration after refinement).

X

X

MMD Threshold t:3.12e-4

MMD: 2.52e-4,NLL:2.29

MMD: 2.17e-4, NLL: 2.20

(b) Results at Iter 4 (middle) and r-Iter 1 (right). MMD and NLL values are shown in the title.

statistic under the null hypothesis (i.e., when dataset X and ˜ X observe the same distribution). The scalar α is the controlled Type-I error (known as the test level), which is set to 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 τ as the empirical ( 1 -α )-quantile 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 by the trained JKO-iFlow model. We use 1000 times of bootstrap in all experiments. In our usage of evaluating generative models, the threshold τ can be viewed as a 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 is below τ , it means that with respect to the current MMD metric, the trained model generates a data distribution that is as good as the true distribution.

C.2 Detail setup and additional results

All experiments are conducted using PyTorch [Paszke et al., 2019] and PyTorch Geometric [Fey and Lenssen, 2019]. The optimizer is Adam [Kingma and Ba, 2015] with learning rates to be specified in each example. We use the neural-ODE integrator [Grathwohl et al., 2019, adjoint method] on all examples.

C.2.1 Two-dimensional toy data

Training and test data: We generate 8K test samples for each dataset, and for the training split,

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

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

Choice of h k : The initial schedule is by setting h 0 = 0 . 75 , ρ = 1 . 2 . We set h max = 5 for rose, checkerboard, and Olympic rings, and h max = 3 for fractal tree. For reparametrization and refinement, we use η = 0 . 5 in the reparameterization iterations for all examples.

For Rose, 3 reparameterization moving iteration, no refinement.

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

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

MMD metric: We use custom kernel bandwidth h = 0 . 1 h m where h m is the median distance of dataset X with N = 8 K (from true data distribution). From trained generative model, we generate M = 10 K test samples as ˜ X .

Network and activation: Fully-connected residual blocks with two hidden layers. We use the softplus activation ( β = 20 ) with 128 hidden dimensions. Before refinement, we train 9 residual blocks for rose, checkerboard, and Olympic rings, and train 6 residual blocks for fractal trees. The learning rate is 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 is 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 high-dimensional real datasets (POWER, GAS, MINIBOONE, BSDS300) come from the University of California Irvine (UCI) machine learning data repository, and we follow the pre-processing procedures of [Papamakarios et al., 2017]. Regarding data sizes,

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

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

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

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

Choice of h k : The initial schedule is by setting h 0 = 1 , ρ = 1 , h max = 3 . For reparametrization and refinement, we use η = 0 . 5 in the reparameterization iterations for all datasets.

MMDmetric: When the test data from the MINIBOONE dataset is less than 10K in size, we use all the test sets as X and N is the size of test data. When the test data from all three other datasets has more than 10K samples, we randomly select N = 10 K 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 = 10 K 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 fully-connected residual blocks with two hidden layers. On POWER, GAS, and MINIBOONE, we use 128 hidden nodes, and on BSDS300, we use 256 hidden nodes. 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-budget 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 in Algorithm 1. (The main lines show the results with L = 4 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 is 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 2 dimensions (determined by principal components computed from true data test samples). The plots show a 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 is used on POWER, GAS, and BSDS300.

(a) Power

(b) Gas

(c) MINIBOONE

(d) BSDS300

TrueX

JKO-iFlow

OT-Flow

FFJORD

IGNN

ScoreSDE

10.0

7.5

5.0

2.5

0'0

2.5

100

05

100

Truex

JKO-iFlow

OT-Flow

FFJORD

IGNN

ScoreSDE

Z

10

True X

JKO-iFlow

OT-Flow

FFJORD

IResNet

ScoreSDE

1

5.0

2.5

7.5

TrueX

JKO-iFlow

OT-Flow

FFJORD

IGNN

ScoreSDE

To ensure a fair comparison against diffusion models, we perform additional experiments using different noise schedulers β ( t ) and ¯ β max in ScoreSDE on MINIBOONE, following the noise scheduler suggestions in DDPM [Ho et al., 2020]. Note that DDPM can be viewed as a discrete-time version of the variance-preserving ScoreSDE model. As shown in Table A.4, the performance of ScoreSDE is indeed sensitive to the noise schedule. However, the best NLL of ScoreSDE 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 is 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 a larger network following the setup in [Albergo and Vanden-Eijnden, 2023], where the network has 4 hidden layers with 512 hidden nodes per layer and consists of 831K parameters in total. We then train this larger model for 100K batches using ScoreSDE with various noise schedulers, choosing noise schedulers that performed best based on results in Table A.4. From the results reported in Table A.5, we see that the testing NLL by ScoreSDE can be as low as 10.47, which is among the state-of-the-art 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 2 hidden layers with 128 hidden nodes per layer, and we trained 4 blocks that contain 112K parameters in total. We also trained JKO-iFlow for only 2.72K batches, rather than the 100K batches we used for ScoreSDE.

C.2.3 Image generation with pre-trained variational auto-encoder

We perform image generation on MNIST [Deng, 2012], CIFAR10 [Krizhevsky and Hinton, 2009], and Imagenet-32 [Deng et al., 2009] datasets. We do so in the latent space of a pre-trained VAE, where we discuss the details below.

VAE as data pre-processing: We train deep VAEs in an adversarial manner following [Esser et al., 2021], and use pre-trained VAEs to pre-process the input images X . Specifically, the encoder E : X → ( µ ( X ) , Σ( X )) of the V AE maps a RGB-image X to parameters of a multivariate Gaussian distribution N ( µ ( X ) , Σ( X )) in a lower dimension ˜ d . Then, given any random latent code ˜ X ∼ N ( µ ( X ) , Σ( X )) , the decoder D : ˜ X → ˆ X of the V AE is trained so that X ≈ ˆ X for the reconstructed image ˆ X = D ( ˜ X ) . On MNIST, we let the latent-space dimension ˜ d = 20 , and on CIFAR10 and ImageNet32, we let the latent-space dimension ˜ d = 192 . There are 5.6M parameters in the VAE encoder and 8.25M parameters in the VAE decoder. Given a trained VAE, we then train the JKO-iFlow model T θ to transport invertibly between the distribution of random latent codes ˜ X defined over all training images X and the standard multivariate Gaussian distribution in R ˜ d .

Training and test data: Training data: 60K images in MNIST, 50K images in CIFAR10, and 1.28M images in Imagenet-32. Test data: 10K images in MNIST and CIFAR10, and 50K images in Imagenet-32.

Choice of h k : on MNIST, we specify { h 0 = 1 , ρ = 1 , h max = 1 } . on CIFAR10, we specify { h 0 = 1 , ρ = 1 , h max = 1 } . On Imagenet-32, we specify { h 0 = 1 , ρ = 1 . 1 , h max = ∞} . We train L = 6 blocks on MNIST, and we train L = 8 blocks on CIFAR10 and Imagenet-32. On CIFAR10 and Imagenet-32, we also scale h k = T · h k / ∑ j h j so ∑ k h k = T . We let T = 0 . 8 on CIFAR10 and T = 0 . 4 on Imagenet-32.

Network architecture: We use the softplus activation with β = 20 for all hidden layers. On MNIST, we use fully connected residual blocks with three hidden layers at 256 hidden nodes. On CIFAR10 and Imagenet-32, we parametrize each f θ b as a concatenation of convolution layers and transposed convolution layers. Specifically:

CIFAR10: convolution layers have channels 3-64-128-256-256 with kernel size 3 and strides 1-1-2-1 . Transposed convolution layers have channels 256-256-128-64-3 with kernel size 3-4-3-3 and strides 1-2-1-1 . Total 2.1M parameters.

Imagenet-32: convolution layers have channels 3-64-128-256-512 with kernel size 3 and strides 1-1-2-1 . Transposed convolution layers have channels 512-256-128-64-3 with kernel size 3-4-3-3 and strides 1-2-1-1 . Total 3.3M parameters.

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 [Ho et al., 2020, Song et al., 2021, Boffi and Vanden-Eijnden, 2023]. For instance, the DDPM model [Ho et al., 2020] on CIFAR10 has 35.7M parameters with more sophisticated model designs.

Training specifics: On MNIST, we fix the batch size to be 2000 during training, and we train 15K batches per JKO-iFlow block. We fix the learning rate to be 1e-3.

On CIFAR10 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 CIFA10. The total time on Imagenet-32 is 30 hours and on CIFAR10 is 24 hours. When training the blocks, the initial learning rate lr0 is decreased by a constant factor of 0.9 every 2500 batches.

12

Initial

Iter 2

Iter 4

Iter6

10

Iter 1

Iter3

Iter 5

Iter 7

8

6

4

2

Q

0

1

2

3

7

(a) Per-block

W

2

2

over reparameterization iterations.

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

(a) Generated CIFAR10 images

X

x=F-1(Z)

人lx

X|Y=F-1(H|Y)

HIY

H|Y=F(X|Y)

(a) Two-circles data (unconditional)

(b) Two-moons data (conditional)

Figure A.7: Unconditional and conditional generation on simulated toy datasets by JKO-iFlow. We color samples in (b) by the class label Y taking binary values. In both (a) and (b), the generated samples are close to the true samples in distribution. In (b), the pushforward distribution by JKO-iFlow, denoted as ˆ H | Y , is also close to the target Gaussian mixture distribution H | Y .

We let lr0 be 1e-3 when training the first 5 blocks on Imagenet-32 and decrease it to be 8e-4 on blocks 6-8. We let lr0 be 1e-3 when training the first 2 blocks on CIFAR10 and decrease it to be 7.5e-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 to curate generated samples for CIFAR10 and Imagenet-32, which are shown in Figure 4, we first train an image classifier (i.e., VGG-16 [Simonyan and Zisserman, 2015]) on the training data. Then, we generate a large number of uncurated images (200K for Imagenet-32 and 20K for CIFAR10), classify them using the pre-trained classifier, and sort them based on top-1 predicted probability. Out of the top 20K images for Imagenet-32 and top 750 images for CIFAR10 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 JKO-iFlow model after training.

C.2.4 Conditional generation

To modify the JKO-iFlow to apply to the conditional generation task, we follow the framework in Xu et al. [2022], which trains a single flow mapping from X to H that pushes to match each component of the distribution associated with a distinct output label value Y = k , k = 1 , · · · , K . Specifically, for a data-label pair { X i , Y i } , consider the continuous ODE trajectory x ( t ) starting from x (0) = X i , the per-sample training objective is by changing the term V ( x ( t k +1 )) in (8) to be V Y i ( x ( t k +1 )) , where V k ( · ) is the potential of the Gaussian mixture component H | Y = k . Because the Gaussian

(a) True

X

|

Y

JKO-iFlow

τ

: 2.13e-3

, MMD-m:

8.12e-3

NLL

-13.60

(b) IGNN

5.70e-2

-4.55

(c) True

X

|

Y

JKO-iFlow

τ

: 2.88e-3,

MMD-m:

4.00e-2

NLL

-4.19

(d) IGNN

7.67e-2

1.53

mixture is parametrized by mean vectors (the covariance is isotropic with fixed variance) [Xu et al., 2022], the expression of V k ( · ) is a quadratic function with explicit expression.

Training and test data:

For the simulated two-moon data, we re-sample 5K training samples every epoch, for a total of 40 epochs per block. The batch size is 1000.

The solar dataset is retrieved from the National Solar Radiation Database (NSRDB), following the data pre-processing in [Xu et al., 2022]. The dataset has a 1K training samples and a 1K test samples. The batch size is 500, and each block is trained for 50 epochs.

Choice of h k : For both datasets, h 0 = 1 , ρ = 1 , h max = 3 . The reparametrization and refinement techniques are not used.

MMDmetric: The MMD values are computed for each specific value of Y (across all graph nodes). When Y is fixed, we retrieve test samples from true data, denoted as X | Y , and generate model distribution samples, denoted as ˜ X | Y , and then compute the MMD values the same as in the unconditional case. Due to the relatively small sample size of X | Y (679 and 144 respectively), we report the average values of MMD-m and threshold τ over 50 replicas. In each replica, we subsample 90% observations of X | Y and generate M = 2000 samples to form ˜ X | Y . For the results in Figure A.8, for the Y in plots (a)(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 is 5e-3.

On solar data, we follow the same design of residual blocks as [Xu 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 11 residual blocks. The learning rate is 5e-3.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nI 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\nWEAKNESSES:\nI 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\nQUESTIONS:\n- 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\nLIMITATIONS:\nNothing to discuss.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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).\n\nQUESTIONS:\nI do not have any particular questions for the authors.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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,\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nCould 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 ?\n\nLIMITATIONS:\nThe authors adequately address the limitations of their work, and there are no obvious potential negative social impacts of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: 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.

STRENGTHS:
1). The proposed method is computationally efficient compared to other normalizing flow methods or diffusion. 

2). Novel idea of optimizing the velocity field instead of the transport map motivated JKO-Scheme. 

3). 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.

2). 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.

QUESTIONS:
1). Is it possible for the authors to provide FID scores for both MNIST and CIFAR-10?

2). How good is the proposed method without using a pre-trained auto-encoder?

LIMITATIONS:
The limitations are clearly discussed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09