SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table A.2: Complete results on real tabular datasets to augment Table 1 under the fixed-budget setting.  $L$  denotes the number of residual blocks used in each dataset; for comparison, the free block in JKO-iFlow is not used. For fair comparison across models, the number of batches indicates how many batches pass through all residual blocks.
Data Set	Model	# Param	Training	Testing
Time (h)	# Batches	Time/Batches (s)	Batch size	MMD-m	MMD-1	NLL
τ: 1.73e-4	τ: 2.90e-4
POWERd=6	JKO-iFlow	76K, L=4	0.07	0.76K	3.51e-1	10000	9.86e-5	2.40e-4	-0.12
OT-Flow	76K	0.36	7.58K	1.71e-1	10000	7.58e-4	5.35e-4	0.32
FFJORD	76K, L=4	0.67	7.58K	3.18e-1	10000	9.89e-4	1.16e-3	0.63
IGNN	304K, L=16	0.29	7.58K	1.38e-1	10000	1.93e-3	1.59e-3	0.95
IResNet	304K, L=16	0.41	7.58K	1.95e-1	10000	3.92e-3	2.43e-2	3.37
ScoreSDE	76K	0.06	7.58K	2.85e-2	10000	9.12e-4	6.08e-3	3.41
ScoreSDE	76K	0.60	75.80K	2.85e-2	10000	7.12e-4	5.04e-3	3.33
JKO-iFlow	57K, L=3	0.05	0.76K	2.63e-1	10000	3.86e-4	7.20e-4	-0.06
τ: 1.85e-4	τ: 2.73e-4
GASd=8	JKO-iFlow	76K, L=4	0.07	0.76K	3.32e-1	5000	1.52e-4	5.00e-4	-7.65
OT-Flow	76K	0.23	7.60K	1.09e-1	5000	1.99e-4	5.16e-4	-6.04
FFJORD	76K, L=4	0.65	7.60K	3.08e-1	5000	1.87e-3	3.28e-3	-2.65
IGNN	304K, L=16	0.34	7.60K	1.61e-1	5000	6.74e-3	1.43e-2	-1.65
IResNet	304K, L=16	0.46	7.60K	2.18e-1	5000	3.20e-3	2.73e-2	-1.17
ScoreSDE	76K	0.03	7.60K	1.42e-2	5000	1.05e-3	8.36e-4	-3.69
ScoreSDE	76K	0.30	76.00K	1.42e-2	5000	2.23e-4	3.38e-4	-5.58
JKO-iFlow	95K, L=5	0.09	0.76K	4.15e-1	5000	1.51e-4	3.77e-4	-7.80
τ: 2.46e-4	τ: 3.75e-4
MINIBOONEd=43	JKO-iFlow	112K, L=4	0.03	0.34K	3.61e-1	2000	9.66e-4	3.79e-4	12.55
OT-Flow	112K	0.21	3.39K	2.23e-1	2000	6.58e-4	3.79e-4	11.44
FFJORD	112K, L=4	0.28	3.39K	2.97e-1	2000	3.51e-3	4.12e-4	23.77
IGNN	448K, L=16	0.63	3.39K	6.69e-1	2000	1.21e-2	4.01e-4	26.45
IResNet	448K, L=16	0.71	3.39K	7.54e-1	2000	2.13e-3	4.16e-4	22.36
ScoreSDE	112K	0.01	3.39K	6.37e-3	2000	5.86e-1	4.33e-4	27.38
ScoreSDE	112K	0.10	33.90K	6.37e-3	2000	4.17e-3	3.87e-4	20.70
τ: 1.38e-4	τ: 1.01e-4
BSDS300d=63	JKO-iFlow	396K, L=4	0.05	1.03K	1.85e-1	1000	2.24e-4	1.91e-4	-153.82
OT-Flow	396K	0.62	10.29K	2.17e-1	1000	5.43e-1	6.49e-1	-104.62
FFJORD	396K, L=4	0.54	10.29K	1.89e-1	1000	5.60e-1	6.76e-1	-37.80
IGNN	990K, L=10	1.71	10.29K	5.98e-1	1000	5.64e-1	6.86e-1	-37.68
IResNet	990K, L=10	2.05	10.29K	7.17e-1	1000	5.50e-1	5.50e-1	-33.11
ScoreSDE	396K	0.01	10.29K	3.50e-3	1000	5.61e-1	6.60e-1	-7.55
ScoreSDE	396K	0.10	102.90K	3.50e-3	1000	5.61e-1	6.62e-1	-7.31
JKO-iFlow	396K, L=4	0.08	1.03K	2.76e-1	5000	1.41e-4	8.83e-5	-156.68
- For Rose, 3 reparameterization moving iteration, no refinement.
- For Checkerboard and Olympic rings, 4 reparameterization moving iteration, no refinement.
- For Fractal tree, 4 reparameterization moving iteration, one refinement, and an additional reparameterization moving iteration.
MMD metric: We use custom kernel bandwidth  $h = 0.1h_{m}$  where  $h_{m}$  is the median distance of dataset  $\mathbf{X}$  with  $N = 8K$  (from true data distribution). From trained generative model, we generate  $M = 10K$  test samples as  $\tilde{\mathbf{X}}$ .
Network and activation: Fully-connected residual blocks with two hidden layers. We use the softplus activation  $(\beta = 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.
Table A.3: The OT-Flow and FFJORD baselines marked with * are from the original papers [Onken et al., 2021] and [Grathwohl et al., 2019]. The results of JKO-iFlow are obtained after applying the reparameterization technique.
Data Set	Model	#Param	Training	Testing NLL
# Batches	Batch size
POWER d = 6	JKO-iFlow	95K, L=5	6.08K	10000	-0.40
OT-Flow*	18K	22K	10000	-0.30
FFJORD*	43K, L=5	-	10000	-0.46
GAS d = 8	JKO-iFlow	114K, L=6	6.08K	5000	-9.43
OT-Flow*	127K	52K	2000	-9.20
FFJORD*	279K, L=5	-	1000	-8.59
MINIBOONE d = 43	JKO-iFlow	112K, L=4	2.72K	2000	10.55
OT-Flow*	78K	7K	2000	10.55
FFJORD*	821K, L=1	-	1000	10.43
BSDS300 d = 63	JKO-iFlow	495K, L=5	2.06K	5000	-157.75
OT-Flow*	297K	37K	300	-154.20
FFJORD*	6.7M, L=2	-	10000	-157.40
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, \rho = 1, h_{\max} = 3$ . For reparametrization and refinement, we use  $\eta = 0.5$  in the reparameterization iterations for all datasets.
MMD metric: When the test data from the MINIBOONE dataset is less than  $10\mathrm{K}$  in size, we use all the test sets as  $\mathbf{X}$  and  $N$  is the size of test data. When the test data from all three other datasets has more than  $10\mathrm{K}$  samples, we randomly select  $N = 10K$  samples from the entire test data, and the same random test subset is used to evaluate all models. We use the median distance kernel bandwidth and generate  $M = 10K$  test samples from each model to form  $\tilde{\mathbf{X}}$ .
Network and activation: We use the softmax 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.
Table A.4: NLL per noise scheduler and  $\bar{\beta}_{\mathrm{max}}$  combination of ScoreSDE on MINIBOONE. The three settings "linear, constant, quadratic" follow the DDPM suggestion [Ho et al., 2020]. The mean and standard deviation are computed over three replicas of the trained model.
Noise scheduler\βmax	1	5	10	15	20
Linear	24.51 (0.24)	18.73 (0.64)	20.28 (0.38)	26.76 (1.77)	26.83 (1.23)
Constant	25.47 (0.25)	30.37 (0.84)	37.73 (0.98)	40.47 (1.27)	45.32 (0.40)
Quadratic	27.19 (0.14)	17.45 (0.17)	18.48 (0.38)	18.90 (0.48)	21.35 (0.60)
Table A.5: Testing NLL per noise scheduler and  $\bar{\beta}_{\mathrm{max}}$  combination of ScoreSDE on MINIBOONE, using a larger network with longer training. The table format is identical to that of Table A.4.
Noise scheduler \ βmax	5	10	15
Linear	11.33 (0.44)	10.84 (1.17)	10.47 (0.13)
Quadratic	12.88 (0.40)	11.11 (0.24)	11.05 (0.49)
(a) Power
(b) Gas
(c) MINIBOONE
(d) BSDS300
Figure A.4: Generative quality on tabular datasets via PCA projection of generated samples. The generative quality in general aligns with the quantitative metrics in Table 1 and A.2.
To ensure a fair comparison against diffusion models, we perform additional experiments using different noise schedulers  $\beta(t)$  and  $\bar{\beta}_{\mathrm{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  $\mathcal{E}: X \to (\mu(X), \Sigma(X))$  of the VAE maps a RGB-image  $X$  to parameters of a multivariate Gaussian distribution  $\mathcal{N}(\mu(X), \Sigma(X))$  in a lower dimension  $\tilde{d}$ . Then, given any random latent code  $\tilde{X} \sim \mathcal{N}(\mu(X), \Sigma(X))$ , the decoder  $\mathcal{D}: \tilde{X} \to \hat{X}$  of the VAE is trained so that  $X \approx \hat{X}$  for the reconstructed image  $\hat{X} = \mathcal{D}(\hat{X})$ . On MNIST, we let the latent-space dimension  $\tilde{d} = 20$ , and on CIFAR10 and ImageNet32, we let the latent-space dimension  $\tilde{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_{\theta}$  to transport invertibly between the distribution of random latent codes  $\tilde{X}$  defined over all training images  $X$  and the standard multivariate Gaussian distribution in  $\mathbb{R}^{\tilde{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, \rho = 1, h_{\max} = 1\}$ . On CIFAR10, we specify  $\{h_0 = 1, \rho = 1, h_{\max} = 1\}$ . OnImagenet-32, we specify  $\{h_0 = 1, \rho = 1.1, h_{\max} = \infty\}$ . We train  $L = 6$  blocks on MNIST, and we train  $L = 8$  blocks on CIFAR10 andImagenet-32. On CIFAR10 andImagenet-32, we also scale  $h_k = T \cdot h_k / \sum_j h_j$  so  $\sum_k h_k = T$ . We let  $T = 0.8$  on CIFAR10 and  $T = 0.4$  onImagenet-32.
Network architecture: We use the softplus activation with  $\beta = 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  $\mathbf{f}_{\theta_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  $1\mathrm{r}0$  is decreased by a constant factor of 0.9 every 2500 batches.
(b) Results at initial training (middle) and Iter 7 (right).
Figure A.5: Reparametrization iterations of JKO-iFlow model on MINIBOONE. (a) After 7 reparameterization iterations, a trajectory of more uniform  $W_{2}$  movements is obtained. (b) Generated samples by the models before and after the moving iterations (visualized in the first two principal components computed from test data).
(a) Generated CIFAR10 images
(b) Generated Imagenet-32 images
Figure A.6: Uncurated generated samples of CIFAR10 and Imagenet-32 by JKO-iFlow in latent space.
(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  $\hat{H} |Y$ , is also close to the target Gaussian mixture distribution  $H|Y$ .
We let 1r0 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 1r0 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, \dots, 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(\cdot)$  is the potential of the Gaussian mixture component  $H|Y = k$ . Because the Gaussian
(a) True  $X|Y$ $\tau$  : 2.13e-3, MMD-m: NLL
(b) IGN 5.70e-2 -4.55
Figure A.8: Conditional graph nodal data generation by JKO-iFlow and iGNN on Solar ramping data. The nodal data  $X \in \mathbb{R}^{10 \times 2}$  and the nodal label  $Y \in \{0,1\}^{10}$ . Plots (a)(b) and (c)(d) are for two different sets of values  $Y$  across all nodes respectively. Samples  $X$  (concatenated across all nodes) are visualized by projection to two principal components (determined by true samples  $X|Y$ ).
JKO-iFlow 8.12e-3 -13.60
(c) True  $X|Y$ $\tau$  : 2.88e-3, MMD-m: NLL
JKO-iFlow (d)IGNN 4.00e-2 7.67e-2 -4.19 1.53
(d) IGN7.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}(\cdot)$  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$ ,  $\rho = 1$ ,  $h_{\max} = 3$ . The reparametrization and refinement techniques are not used.
MMD metric: 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  $\tilde{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  $\tau$  over 50 replicas. In each replica, we subsample  $90\%$  observations of  $X|Y$  and generate  $M = 2000$  samples to form  $\tilde{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.