SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $x(h) = x(0) + \int_0^h\mathbf{f}(x(s),s)ds$ . In the limit of  $h\to 0+$ , formally,  $x(h) - x(0) = h\mathbf{f}(x(0),0) + O(h^{2})$ , and suppose  $V$  of  $q$  is  $C^2$ ,  $V(x(h)) = V(x(0)) + h\nabla V(x(0))\cdot \mathbf{f}(x(0),0) + O(h^2)$ . For any differentiable density  $\rho$ , the (Stein) score function is defined as  $\mathbf{s}_{\rho} = \nabla \log \rho$ , and we have  $\nabla V = -\mathbf{s}_q$ . Taking the formal expansion of orders of  $h$ , the objective in (9) is written as
$$
\mathbb {E} _ {x \sim p} \left(V (x) + h \left(- \mathbf {s} _ {q} (x) \cdot \mathbf {f} (x, 0) - \nabla \cdot \mathbf {f} (x, 0) + \frac {1}{2} \| \mathbf {f} (x, 0) \| ^ {2}\right) + O \left(h ^ {2}\right)\right). \tag {10}
$$
Note that  $\mathbb{E}_{x\sim p}V(x)$  is independent of  $\mathbf{f}(x,t)$ , and the  $O(h)$  order term in (10) is over  $\mathbf{f}(x,0)$  only, thus the minimization of the leading term is equivalent to
$$
\min  _ {\mathbf {f} (\cdot) = \mathbf {f} (\cdot , 0)} \mathbb {E} _ {x \sim p} \left(- T _ {q} \mathbf {f} + \frac {1}{2} \| \mathbf {f} \| ^ {2}\right), \quad T _ {q} \mathbf {f} := \mathbf {s} _ {q} \cdot \mathbf {f} + \nabla \cdot \mathbf {f}, \tag {11}
$$
where  $T_{q}$  is known as the Stein operator [Stein, 1972]. The  $T_{q}\mathbf{f}$  in (11) echoes that the derivative of KL divergence with respect to transport map gives Stein operator [Liu and Wang, 2016]. The Wasserstein-2 regularization gives an  $L^2$  regularization in (11). Let  $L^{2}(p)$  be the  $L^2$  space on  $(\mathbb{R}^d,p(x)dx)$ , and for vector field  $\mathbf{v}$  on  $\mathbb{R}^d$ ,  $\mathbf{v}\in L^{2}(p)$  if  $\int |\mathbf{v}(x)|^2 p(x)dx < \infty$ . One can verify that, when both  $\mathbf{s}_p$  and  $\mathbf{s}_q$  are in  $L^2 (p)$ , the minimizer of (11) is
$$
\mathbf {f} ^ {*} (\cdot , 0) = \mathbf {s} _ {q} - \mathbf {s} _ {p}.
$$
This shows that the infinitesimal optimal  $\mathbf{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  $\rho_{t}$ , the analysis in Section 3.2 implies that the optimal vector field  $\mathbf{f}(x,t)$  has the expression as
$$
\mathbf {f} (x, t) = \mathbf {s} _ {q} - \mathbf {s} _ {\rho_ {t}} = - \nabla V - \nabla \log \rho_ {t}. \tag {12}
$$
With this  $\mathbf{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  $\rho_0$ , the r.h.s. of (12) is also regular over space and time. This leads to two consequences, in approximation and in learning: Approximation-wise, the regularity of  $\mathbf{f}(x,t)$  allows to construct a  $k$ -th residual block in the flow network to approximate  $\{\mathbf{f}(x,t)\}_{t\in I_k}$  when there is sufficient model capacity, by classical universal approximation theory of shallow networks [Barron, 1993, Yarotsky, 2017]. 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  $\mathbf{f}(x,t;\theta_k)$  will learn to approximate (12). This can be viewed as inferring the score function of  $\rho_{t}$ , and also leads to the invertibility of the trained flow net in theory: Suppose the trained  $\mathbf{f}(x,t;\theta_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  $\mathbf{f}(x,t;\theta_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)\}_{i = 1}^{n}$ , the expectation  $\mathbb{E}_{x(t)\sim p_k}$  in (8) is replaced by the sample average over  $\{x_{i}(kh)\}_{i = 1}^{n}$  which can be computed from the previous  $k$  blocks. Note that for each given  $x(t) = x(t_{k})$ , both  $x(t_{k + 1})$  and the integral of  $\nabla \cdot \mathbf{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  $\nabla \cdot \mathbf{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_{\theta_k}$  with trainable parameters  $\theta_{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 \coloneqq 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  $\mathrm{Ter}(k)$  in line 2 of Algorithm 1, we monitor the ratio  $\mathbb{E}_{x \sim p_{k-1}} \| x - T_k(x) \|^2 / \mathbb{E}_{x \sim p_{k-1}} \| T_k(x) \|^2$  and terminate when it is below some threshold  $\epsilon$ , set as 0.01 in all experiments. 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.,
by Runge-Kutta-4. We empirically verify small inversion errors on test samples.
Algorithm 1 Block-wise JKO-iFlow training
Require: Time stamps  $\{t_k\}$ , training data, termination criterion Ter and tolerance level  $\epsilon$ , maximal number of blocks  $L_{\mathrm{max}}$ .
1: Initialize  $k = 1$ .
2: while  $\operatorname{Ter}(\mathrm{k}) > \epsilon$  and  $k \leq L_{\max}$  do
3: Optimize  $f_{\theta_k}$  upon minimizing (8) with mini-batch sample approximation, given  $\{f_{\theta_i}\}_{i=1}^{k-1}$ . Set  $k \gets k + 1$ .
4: end while
5:  $L\gets k$  . Optional: Optimize  $f_{\theta_{L + 1}}$  without  $W_{2}$  regularization.
4.2 Computation of trajectories in probability space
We adopt two additional computational techniques to facilitate learning of the trajectories in the probability space, represented by the sequence of densities  $p_k$ ,  $k = 1, \dots, 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
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.
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  $\tau$  is computed by bootstrap, where an MMD metric less than  $\tau$  indicates that the generated distribution is evaluated by the MMD test to be the same as the true data distribution (achieving  $\alpha = 0.05$  test level). See details in Appendix C.1. The computational cost is measured by the number of mini-batch stochastic gradient descent steps and the training
JKO-iFlow
(b)FFJORD
(c) OT-Flow
(d)IGNN
(e) ScoreSDE
$\tau$  : 2.79e-4. MMD-c:
NLL 
(f) Fractal tree
$\tau$  :3.12e-4.MMD-c:
NLL
Figure 3: Results on two-dimensional simulated datasets by JKO-iFlow and competitors.
2.73e-4
2.64
3.88e-4
2.95
2.17e-4
2.20
(g) Olympic rings
$\tau$  :3.16e-4 MMD-c:
NLL
1.42e-3
3.30
2.36e-4
1.66
3.14e-3
3.35
(h) Checkerboard
T:3.09e-4.MMD-c:
NLL
6.90e-4
3.2
2.70e-4
3.59
Table 1: Results on tabular datasets. All competitors are trained in a fixed-budget setup using 10 times more mini-batches (their performances using the same number of mini-batches are worse and not comparable to JKO-iFlow). See the complete table in Table A.2. The results using more computation are given in Table A.3.
Data Set	Model	# Param	Test MMD-m	Test MMD-1	NLL
τ: 1.73e-4	τ: 2.90e-4
POWDERd=6	JKO-iFlow	76K	9.86e-5	2.40e-4	-0.12
OT-Flow	76K	7.58e-4	5.35e-4	0.32
FFJORD	76K	9.89e-4	1.16e-3	0.63
IGNNN	304K	1.93e-3	1.59e-3	0.95
IResNet	304K	3.92e-3	2.43e-2	3.37
ScoreSDE	76K	9.12e-4	6.08e-3	3.41
τ: 1.85e-4	τ: 2.73e-4
GASd=8	JKO-iFlow	76K	1.52e-4	5.00e-4	-7.65
OT-Flow	76K	1.99e-4	5.16e-4	-6.04
FFJORD	76K	1.87e-3	3.28e-3	-2.65
IGNNN	304K	6.74e-3	1.43e-2	-1.65
IResNet	304K	3.20e-3	2.73e-2	-1.17
ScoreSDE	76K	1.05e-3	8.36e-4	-3.69
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  $\hat{X}$  by JKO-iFlow also more resemble the true data distribution of  $X$  than the competitors in (b)-(e). Plots (f)-(h) show the results by JKO-iFlow on the examples of Fractal tree, Olympic rings, and Checkerboard, where 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 Imagenet-32 datasets. In all examples, we train JKO-iFlow in the latent space of a pre-trained variational auto-encoder (VAE) adopted from [Esser et al., 2021]. Training generative models in latent space 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échet inception distance (FID) score [Heusel et al., 2017], which is included in the figure captions. Comparatively, we are aware of some performance gap between our model and the state-of-the-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 onImagenet-32. Meanwhile, the image generation by JKO-iFlow model here obtains visually more appealing images and achieves lower
(a) Generated MNIST digits. FID: 7.95.
(b) Generated CIFAR10 images. FID: 29.10.
(c) Generated Imagenet-32 images. FID: 20.10.
Figure 4: Generated samples of MNIST, CIFAR10, and Imagenet-32 by JKO-iFlow model in latent space. We select 2 images per class for CIFAR10 and 1 image per class for Imagenet-32. The FIDs are shown in subcaptions. Uncurated samples are shown in Figure A.6.
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.
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  $\mathbf{f}(x,t)$  as in (12) can lead to approximation of the Fokker-Planck flow (3), which pushes forward the density to be  $\varepsilon$ -close to normality in  $T = \log(1 / \varepsilon)$  time, see (4). Reversing the time of the ODE then leads to an approximation of the initial density  $\rho_0 = p_X$  by flowing backward in time from  $T$  to zero. Further analysis under technical assumptions is left to future work. During the time that this paper was being published, the convergence analysis of the proposed model was studied in Cheng et al. [2023].
The authors thank Jianfeng Lu, Yulong Lu, and Yiping Lu for helpful discussions. The work is partially supported by NSF DMS-2134037. C.X. and Y.X. are partially supported by an NSF CAREER CCF-1650913, and NSF DMS-2134037, CMMI-2015787, CMMI-2112533, DMS-1938106, and DMS-1830210 and the Coca-Cola Foundation. XC is also partially supported by NSF DMS-2237842 and Simons Foundation.
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örn-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.
François Bolley, Ivan Gentil, and Arnaud Guillin. Convergence to equilibrium in Wasserstein distance for fokker-planck equations. Journal of Functional Analysis, 263(8):2430-2457, 2012.
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örn-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ël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. Advances in neural information processing systems, 29, 2016.