SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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ère 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.1 Proofs in Section
Lemma A.1. Suppose  $p$  and  $q$  are two densities on  $\mathbb{R}^d$  in  $\mathcal{P}$ , the following two problems
$$
\min  _ {\rho \in \mathcal {P}} L _ {\rho} [ \rho ] = \mathrm {K L} (\rho | | q) + \frac {1}{2 h} W _ {2} ^ {2} (p, \rho), \tag {13}
$$
$$
\min  _ {T: \mathbb {R} ^ {d} \rightarrow \mathbb {R} ^ {d}} L _ {T} [ T ] = \mathrm {K L} (T _ {\#} p \| q) + \frac {1}{2 h} \mathbb {E} _ {x \sim p} \| x - T (x) \| ^ {2}, \tag {14}
$$
have the same minimum and
(a) If  $T^{*} : \mathbb{R}^{d} \to \mathbb{R}^{d}$  is a minimizer of (14), then  $\rho^{*} = (T^{*})_{\#} p$  is a minimizer of (13).
(b) If  $\rho^{*}$  is a minimizer of (13), then the optimal transport from  $p$  to  $\rho^{*}$  minimizes (14).
Proof of Lemma A.1. Let the minimum of (14) be  $L_T^*$ , and that of (13) be  $L_{\rho}^{*}$ .
Proof of (a): Suppose  $L_{T}$  achieves minimum at  $T^{*}$ , then  $T^{*}$  is the optimal transport from  $p$  to  $\rho^{*} = (T^{*})_{\#}p$  because otherwise  $L_{T}$  can be further improved. By definition of  $L_{\rho}$ , we have  $L_{T}^{*} = L_{T}[T^{*}] = L_{\rho}[\rho^{*}] \geq L_{\rho}^{*}$ . We claim that  $L_{T}^{*} = L_{\rho}^{*}$ . Otherwise, there is another  $\rho'$  such that  $L_{\rho}[\rho'] < L_{T}^{*}$ . Let  $T'$  be the optimal transport from  $p$  to  $\rho'$ , and then  $L_{T}[T'] = L_{\rho}[\rho'] < L_{T}^{*}$ , contradicting with that  $L_{T}^{*}$  is the minimum of  $L_{T}$ . This also shows that  $L_{\rho}[\rho^{*}] = L_{T}^{*} = L_{\rho}^{*}$ , that is,  $\rho^{*}$  is a minimizer of  $L_{\rho}$ .
Proof of (b): Suppose  $L_{\rho}$  achieves minimum at  $\rho^{*}$ . Let  $T^{*}$  be the OT from  $p$  to  $\rho^{*}$ , then  $\mathbb{E}_{x \sim p}|x - T^{*}(x)|^{2} = W_{2}(p, \rho^{*})^{2}$ , and then  $L_{T}[T^{*}] = L_{\rho}[\rho^{*}] = L_{\rho}^{*}$  which equals  $L_{T}^{*}$  as proved in (a). This shows that  $T^{*}$  is a minimizer of  $L_{T}$ .
Proof of Proposition 3.1. Given  $p_k$  being the density of  $x(t)$  at  $t = kh$ , recall that  $T$  is the solution map from  $x(t)$  to  $x(t + h)$ . We denote  $\rho_t \coloneqq p_k$ , and  $\rho_{t + h} \coloneqq T_{\#}p_k$ . By definition,
$$
\operatorname {K L} \left(T _ {\#} p _ {k} \| p _ {Z}\right) = \mathbb {E} _ {x \sim \rho_ {t + h}} \left(\log \rho_ {t + h} (x) - \log p _ {Z} (x)\right). \tag {15}
$$
Because  $p_Z \propto e^{-V}$ ,  $V(x) = |x|^2 / 2$ , we have  $\log p_Z(x) = -V(x) + c_1$  for some constant  $c_1$ . Thus
$$
\mathbb {E} _ {x \sim \rho_ {t + h}} \log p _ {Z} (x) = \mathbb {E} _ {x (t) \sim \rho_ {t}} \log p _ {Z} (x (t + h)) = c _ {1} - \mathbb {E} _ {x (t) \sim \rho_ {t}} V (x (t + h)). \tag {16}
$$
To compute the first term in (15), note that
$$
\mathbb {E} _ {x \sim \rho_ {t + h}} \log \rho_ {t + h} (x) = \mathbb {E} _ {x (t) \sim \rho_ {t}} \log \rho_ {t + h} (x (t + h)), \tag {17}
$$
and by the expression (called "instantaneous change-of-variable formula" in normalizing flow literature [Chen et al., 2018], which we derive directly below)
$$
\frac {d}{d t} \log \rho (x (t), t) = - \nabla \cdot \mathbf {f} (x (t), t), \tag {18}
$$
we have that for each value of  $x(t)$
$$
\log \rho_ {t + h} (x (t + h)) = \log \rho (x (t + h), t + h) = \log \rho (x (t), t) - \int_ {t} ^ {t + h} \nabla \cdot \mathbf {f} (x (s), s) d s.
$$
Inserting back to (17), we have
$$
\mathbb {E} _ {x \sim \rho_ {t + h}} \log \rho_ {t + h} (x) = \mathbb {E} _ {x (t) \sim \rho_ {t}} \log \rho_ {t} (x (t)) - \mathbb {E} _ {x (t) \sim \rho_ {t}} \int_ {t} ^ {t + h} \nabla \cdot \mathbf {f} (x (s), s) d s).
$$
The first term is determined by  $\rho_{t} = p_{k}$ , and thus is a constant  $c_{2}$  independent from  $\mathbf{f}(x,t)$  on  $t\in [kh,(k + 1)h]$ . Together with (16), we have shown that
$$
\mathrm {r . h . s . o f} (1 5) = c _ {2} - \mathbb {E} _ {x (t) \sim \rho_ {t}} \int_ {t} ^ {t + h} \nabla \cdot \mathbf {f} (x (s), s) d s) - c _ {1} + \mathbb {E} _ {x (t) \sim \rho_ {t}} V (x (t + h)),
$$
which proves (7).
Derivation of (18): by chain rule,
$$
\begin{array}{l} \frac {d}{d t} \log \rho (x (t), t) = \frac {\nabla \rho (x (t) , t) \cdot \dot {x} (t) + \partial_ {t} \rho (x (t) , t)}{\rho (x (t) , t)} \\ = \left. \frac {\nabla \rho \cdot \mathbf {f} - \nabla \cdot (\rho \mathbf {f})}{\rho} \right| _ {(x (t), t)} \quad (\text {b y} (1) \text {a n d} (2)) \\ = - \nabla \cdot \mathbf {f} (x (t), t). \\ \end{array}
$$
B Technical details of Section 4
To train the proposed JKO-iFlow model by Algorithm 1, one needs to specify a sequence of  $t_k$ , where the  $k$ -th JKO block integrates from  $t_{k-1}$  to  $t_k$ , starting from  $t_0 = 0$ . Denote by  $h_k := t_{k+1} - t_k$  the step-size for  $k$ -th step.
B.1 Initial choice of  $t_k$  in Algorithm 1
We first initialize the choice by the following scheme
$$
h _ {k} = \min  \left\{\rho^ {k} h _ {0}, h _ {\max } \right\}, \quad k = 0, 1, \dots \tag {19}
$$
where the base stepsize  $h_0$ , the multiplying factor  $\rho \geq 1$ , and the maximum stepsize  $h_{\mathrm{max}} \geq h_0$  are three hyper-parameters to be specified by the user. When  $\rho = 1$ , the sequence of  $h_k$  in (19) corresponds to constant stepsize  $h_k \equiv h_0$ .
Applying Algorithm 1 with the initial choice of  $t_k$  gives  $L$  trained residual blocks of the neural ODE model, where  $L$  is to be determined by the stopping criterion in Section 4.1. The model is further improved by the two additional techniques in Section 4.2, which will adaptively specify the stepsize  $h_k$  and we explain in detail below.
B.2 Hutchinson trace via finite difference
We propose a finite-difference estimator of  $\nabla \cdot \mathbf{f}(x,t) = \mathrm{Tr}(J_{\mathbf{f}}(x))$ , where  $J_{\mathbf{f}}(x)$  is the Jacobian of  $\mathbf{f}$  at  $x$ . This is based on the Hutchinson trace estimator [Hutchinson, 1989], which states that  $\mathrm{Tr}(J_{\mathbf{f}}(x)) = \mathbb{E}_{p(\epsilon)}\left[\epsilon^T J_{\mathbf{f}}(x)\epsilon \right]$ . Here,  $p(\epsilon)$  is a distribution in  $\mathbb{R}^d$  satisfying  $\mathbb{E}[\epsilon] = 0$  and  $\mathrm{Cov}(\epsilon) = I$  (e.g., a standard Gaussian). Existing CNF works estimate the expectation by sampling  $\epsilon \sim p(\epsilon)$  and computing the vector-Jacobian product  $J_{\mathbf{f}}(x)\epsilon$  using reverse-mode automatic differentiation. The product  $J_{\mathbf{f}}(x)\epsilon$  can be computed for approximately the same cost as evaluating  $\mathbf{f}$  [Grathwohl et al., 2019].
Now, given a fixed  $\epsilon$ , we have  $J_{\mathbf{f}}(x)\epsilon = \lim_{\sigma \to 0}\frac{\mathbf{f}(x + \sigma\epsilon) - \mathbf{f}(x)}{\sigma}$ , which is the directional derivative of  $\mathbf{f}$  along the direction  $\epsilon$ . Thus, given a fixed small  $\sigma > 0$ , we propose the following finite-difference estimator of  $\nabla \cdot \mathbf{f}(x,t)$ :
$$
\nabla \cdot \mathbf {f} (x, t) \approx \mathbb {E} _ {p (\epsilon)} \left[ \epsilon^ {T} \frac {\mathbf {f} (x + \sigma \epsilon) - \mathbf {f} (x)}{\sigma} \right]. \tag {20}
$$
The finite-difference approximation of  $J_{\mathbf{f}}(x)\epsilon$  requires only one additional function evaluation of  $\mathbf{f}$  at  $x + \sigma \epsilon$  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\sim 2$  times faster than existing CNF approaches relying on reverse-mode automatic differentiation. In all experiments, we let  $\sigma = \sigma_0 / \sqrt{d}$  with  $\sigma_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  $\mathbb{R}^d$ , and  $F: \mathbb{R}^d \to \mathbb{R}$  be a differential landscape. Given a sequence of positive stepsize  $h_k$ , one can compute a sequence of representative points  $x_k$  which discretizes the gradient descent trajectory of minimizing  $F$  (towards a local minimum  $x^*$ ). Specifically, we have
$$
x _ {k + 1} = \arg \min  _ {x} F (x) + \frac {1}{2 h _ {k}} \| x - x _ {k} \| _ {2} ^ {2}, \tag {21}
$$
starting from some fixed  $x_0$ .
B.3.1 Trajectory reparameterization
Starting from an initial sequence of step-size  $\{h_k\}_{k=1}^L$  as in Section B.1, the reparameterization technique implements an iterative scheme to update the sequence of  $h_k$  adaptively. We call the  $j$ -th iteration 'Iter- $j$ , and denote the sequence by  $\{h_k^{(j)}\}_{k=1}^L$ . The corresponding solution points  $x_k$  by solving (21) with  $h_k^{(j)}$  are denoted as  $x_k^{(j)}$ . In addition to the constant  $h_{\max}$ , we need another algorithmic parameter  $\eta \in (0,1)$ , and we set  $\eta = 0.3$  for the vector-space example in computation.
We always have  $x_0^{(j)} = x_0$  for all 'Iter- $j$ . In the  $j$ -th iteration of the parametrization, we compute the following
1. Compute the arclength as
$$
S _ {k} ^ {(j)} := \left\| x _ {k - 1} ^ {(j)} - x _ {k} ^ {(j)} \right\| _ {2}, \quad k = 1, \dots , L. \tag {22}
$$
2. Compute the average arclength
$$
\bar {S} := \sum_ {k = 1} ^ {L} S _ {k} ^ {(j)} / L.
$$
3. Update the stepsize for  $k = 1, \dots, L$  as
$$
h _ {k} ^ {(j + 1)} := \min  \left\{h _ {k} ^ {(j)} + \eta (\bar {S} h _ {k} ^ {(j)} / S _ {k} ^ {(j)} - h _ {k} ^ {(j)}), h _ {\max } \right\}
$$
4. Solve (21) using  $h_k^{(j + 1)}$  to obtain  $x_{k}^{(j + 1)}$ .
We terminate the iteration when the arclength  $\{S_k^{(j)}\}$  are approximately equal. An illustration is given as Iter-12 in the upper panel of Figure A.1. After the reparameterization iterations, we solve an additional  $x_{L + 1}$  by minimizing  $F(x)$  starting from  $x_{L}$ . This corresponds to optimizing the "free-block" in Algorithm 1.
B.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\}_{k=1}^L$  and the representative points  $\{x_k\}_{k=1}^L$ , the refinement scheme computes a refined trajectory having  $k = 1, \dots, 2L$ :
1. Compute the refined step sizes as
$$
\tilde {h} _ {2 k - 1} = \tilde {h} _ {2 k} = h _ {k} / 2, \quad k = 1, \dots , L
$$
2. Compute the representative points  $\tilde{x}_k$  by solving (21) using  $\tilde{h}_k$ , possibly warm-starting by  $\tilde{x}_{2k} = x_k$  and  $\tilde{x}_{2k-1} = (x_k + x_{k-1}) / 2$ .
We then apply the trajectory reparameterization iterations as in Section B.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  $\mathcal{P}$ , and the optimization is over the parametrization of  $p_{k}$  induced by the neural network flow mapping consisting of the first  $k$  residual blocks.
It remains to define the arclength in (22) to implement the two techniques. Because we equip  $\mathcal{P}$  with the Wasserstein-2 metric, we compute (omitting the iteration index  $j$  in the notation)
$$
S _ {k} = W _ {2} \left(p _ {k - 1}, p _ {k}\right) = \left(\mathbb {E} _ {x \sim p _ {k - 1}} \| x - T _ {k} (x) \| ^ {2}\right) ^ {1 / 2}, \tag {23}
$$
where the transport mapping  $T_{k}(x) = x + \int_{t_{k - 1}}^{t_{k}}f_{\theta_{k}}(x(s),s)ds$  and can be computed from the  $k$ -th block. In practice, the expectation  $\mathbb{E}_{x\sim p_{k - 1}}$  in (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  $\theta_k$  of the trained blocks.
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.
Figure A.1: The upper panel shows the arc lengths and the mean and standard deviation of arc lengths over 12 reparameterization iterations, one refinement, and an additional 10 reparameterization iterations. The lower panel visualizes the trajectory consisting of 17 solution points at "r-Iter-10", where the free point as the 17-th point computes the approximate minimum by  $x_{17} = x_{L + 1} = [-1.911, 0.105]$ .
C.1.1 NLL metric
Our computation of the NLL metric follows the standard procedure for neural ODE normalizing flow models Chen et al. [2018], Grathwohl et al. [2019], where the evolution of density  $\rho(x(t))$  can be computed via integrating  $\nabla \cdot \mathbf{f}(x(t), t)$  over time due to the instantaneous change-of-variable formula (18).
Specifically, our model flows from  $t_0 = 0$  to  $t_{L + 1} = T$ , where  $x(0) \sim p_0$  the data distribution, and we train the flow model to make  $x(T)$  follow a normal density  $q$ . The model density at data sample  $x$  is expressed as  $\rho(x,0) = (T_\theta^{-1})_\# q(x)$  where  $T_\theta^{-1}$  is the inverse model flow mapping from  $x(T)$  to  $x(0)$ . Thus the model log-likelihood can be expressed as
$$
\log \rho (x (0), 0) = \log q (x (T), T) + \int_ {0} ^ {T} \nabla \cdot \mathbf {f} (x (s), s) d s.
$$
In practice, our trained flow model has  $L$  blocks where each block  $f_{\theta_k}$  represents  $\mathbf{f}(x,t)$  on  $[t_{k - 1},t_k)$  and is parametrized by  $\theta_{k}$ . The log-likelihood at test sample  $x$  is then computed by
$$
\operatorname {L L} (x) = - \frac {1}{2} \left(\| x \left(t _ {L + 1}\right) \| _ {2} ^ {2} + d \log (2 \pi)\right) + \sum_ {k = 1} ^ {L + 1} \int_ {t _ {k - 1}} ^ {t _ {k}} \nabla \cdot f _ {\theta_ {k}} (x (s), s) d s, \tag {24}
$$
$$
x \left(t _ {k}\right) = x \left(t _ {k - 1}\right) + \int_ {t _ {k - 1}} ^ {t _ {k}} f _ {\theta_ {k}} (x (s), s) d s
$$
starting from  $x(0) = x$ . Both the integration of  $f_{\theta_k}$  and  $\nabla \cdot f_{\theta_k}$  are computed using the numerical scheme of neural ODE. We report NLL in the natural unit of information (i.e., log with base  $e$ , known as "nats") in all our experiments.
C.1.2 MMD metrics
Note that the normalizing flow models use NLL (on training set) as the training objective. In contrast, the MMD metric is an impartial evaluation metric as it is not used to train JKO-iFlow or any competing methods. Given two sets of data samples  $\mathbf{X} \coloneqq \{x_{i}\}_{i = 1}^{N}$  and  $\tilde{\mathbf{X}} \coloneqq \{\tilde{x}_j\}_{j = 1}^M$  and a kernel function  $k(x,\tilde{x})$ , the (squared) kernel MMD [Gretton et al., 2012a] is defined as
$$
\operatorname {M M D} (\boldsymbol {X}, \tilde {\boldsymbol {X}}) := \frac {1}{N ^ {2}} \sum_ {i = 1} ^ {N} \sum_ {j = 1} ^ {N} k \left(x _ {i}, x _ {j}\right) + \frac {1}{M ^ {2}} \sum_ {i = 1} ^ {M} \sum_ {j = 1} ^ {M} k \left(\tilde {x} _ {i}, \tilde {x} _ {j}\right) - \frac {2}{N M} \sum_ {i = 1} ^ {N} \sum_ {j = 1} ^ {M} k \left(x _ {i}, \tilde {x} _ {j}\right), \tag {25}
$$
When a generative model is trained, we generate  $M$  i.i.d. data samples by the model to construct the set  $\tilde{\pmb{X}}$ , and we form the set  $\pmb{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,\tilde{x}) = \exp \{-\| x - \tilde{x}\|^{2} / 2h^{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 = 10K$  test samples as  $\mathbf{X}$ , and for each trained model, we generate  $M = 10K$  test samples to form the dataset  $\tilde{\mathbf{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  $\tau$ , which is defined as the upper  $(1 - \alpha)$ -quantile of the distribution of the MMD
(a) Per-block  $W_{2}^{2}$  over reparameterization iterations.
(b) Results at initial training (middle) and Iter 3 (right). MMD and NLL values are shown in the title.
Figure A.2: Same plots as in Figure A.5 for Rose data. We observe improved generative quality on the leaves of the rose after reparameterization iterations
Figure A.3: Same plots as in Figure A.5 for Fractal tree data. We observe improved generative quality on the edges of the tree after refinement.
(a) Per-block  $W_{2}^{2}$  over reparameterization iterations and refinement ('r-Iter 1' means one reparameterization iteration after refinement).
(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  $\mathbf{X}$  and  $\tilde{\mathbf{X}}$  observe the same distribution). The scalar  $\alpha$  is the controlled Type-I error (known as the test level), which is set to be 0.05. To obtain the test threshold  $\tau$  for the MMD, we adopt the bootstrap procedure Arcones and Gine [1992], Gretton et al. [2012a] and compute  $\tau$  as the empirical  $(1 - \alpha)$ -quantile of the simulated null distribution of the MMD from the pool of samples formed by the union of  $\mathbf{X}$  and  $\tilde{\mathbf{X}}$ , where the set  $\tilde{\mathbf{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  $\tau$  can be viewed as a baseline of the MMD metric, that is, when the computed MMD values are above  $\tau$ , then the smaller the MMD value the better the generative performance; when the computed MMD value is below  $\tau$ , it means that with respect to the current MMD metric, the trained model generates a data distribution that is as good as the true distribution.
C.2 Detail setup and additional results
All experiments are conducted using PyTorch [Paszke et al., 2019] and PyTorch Geometric [Fey and Lenssen, 2019]. The optimizer is Adam [Kingma and Ba, 2015] with 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, \rho = 1.2$ . We set  $h_{\mathrm{max}} = 5$  for rose, checkerboard, and Olympic rings, and  $h_{\mathrm{max}} = 3$  for fractal tree. For reparametrization and refinement, we use  $\eta = 0.5$  in the reparameterization iterations for all examples.
Table A.1: Inversion error  ${\mathbb{E}}_{x \sim  {p}_{X}}\begin{Vmatrix}{T}_{\theta }^{-1}\left( {{T}_{\theta }\left( x\right) }\right)  - x\end{Vmatrix}^{2}$  of JKO-iFlow computed via sample average on the test split of the data set, where  ${T}_{\theta }$  denotes the transport mapping over all the blocks of the trained flow network.
Rose	Fractal tree	Olympic rings	Checkerboard	POWER	GAS	MINIBOONE	BSD300	MNIST
3.30e-6	3.58e-5	2.24e-6	3.07e-5	1.48e-5	1.58e-6	1.09e-6	1.53e-5	1.87e-5