SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For our mixture benchmark pairs experiment, we adapt the author's setup for the Gaussian-to-Gaussian experiment from their original paper [42, §5.2]. In particular, we use the same architectures of neural networks, see [42, Appendix C.2], but change the hyper-parameters of [42, Algorithm 1], since the orig
Table 8: | NOT| training parameters
K	Ktest	√η	σ0	N
500	1000	0.05	1.0	1024
Table 9: [EgNOT] training parameters for the mixture benchmark pairs experiment.
inal ones do not work properly when fitting Gaussian-to-Mixture. We hypothesize that the observed failure is due to the short-run nature of the energy-based training algorithm. We suppose that significantly increasing the number of Langevin steps  $K$  used at the training stage may leverage the problem. The specific hyper-parameters of  $\lfloor \mathbf{EgNOT}\rfloor$  algorithm are the same for all  $(\epsilon ,D)$  pairs and provided in Table 9.
We initialize the learning rate as  $lr = 10^{-5}$  and decrease its value during the training. Similar to the original implementation of [42] we use a replay buffer but found that a high probability ( $p = 0.95$ ) of samples reuse does not improve the quality and sometimes leads to unstable training. In turn, we choose  $p = 0.5$ . The reported numbers in Tables 4, 5 are gathered by launching the training process
for approximately 50K iterations and reporting the best-obtained metric. We understand that such an evaluation procedure is not ideal and does not provide statistically significant results. However, the qualitative results reported in Table 2 seem to show the behaviour of [EgNOT] solver on our benchmark setup and reveal the key properties of the approach.
[ENOT] [23] We use the official code from
https://github.com/ngushchin/EntropicNeuralOptimalTransport
We use the same hyperparameters for this setup as the authors [23, Appendix E], except the number of discretization steps  $\mathbf{N}$ , which we set to 200 as well as for other Schrödinger Bridge based methods. We also change the learning rate of the potential to  $3 \cdot 10^{-4}$  for the setups with  $\epsilon = 10$ .
[MLE-SB] [52]. We tested the official code from
https://github.com/franciscovargas/GP_Sinkhorn
Instead of Gaussian processes, we used a neural network as for  $\lfloor \mathbf{ENOT}\rfloor$ . We use  $N = 200$  discretization steps as for other SB solvers, 5000 IPF iterations, and 512 samples from distributions  $\mathbb{P}_0$  and  $\mathbb{P}_1$  in each of them. We use the Adam optimizer with  $lr = 10^{-4}$  for optimization.
DiffSB[15]. We utilize the official code from
https://github.com/JTT94/diffusion_schrodinger_bridge
with their configuration blob/main/conf/dataset/2d.yaml for toy problems. We increase the number of steps of dynamics to 200 and the number of steps of the IPF procedure for dimensions 16, 64 and 128 to 30, 40 and 60, respectively.
$\lfloor \mathbf{FB - SDE - J}\rfloor [9]$  .We utilize the official code from
https://github.com/ghliu/SB-FBSDE
with their configuration blob/main/configs/defaultcheckerboard_config.py for the checkerboard-to-noise toy experiment, changing the number of steps of dynamics from 100 to 200 steps. Since their hyper-parameters are developed for their 2-dimensional experiments, we increase the number of iterations for dimensions 16, 64 and 128 to 15000.
[FB-SDE-A] [9]. We also take the code from the same repository as above. We base our configuration on the authors' one (blob/main/configs/defaultmoon_to_spiral_config.py) for the moon-to-spiral experiment. As earlier, we increase the number of steps of dynamics up to 200. Also, we change the number of training epochs during one IPF procedure for dimensions 16, 64 and 128 to 2,4 and 8 correspondingly.
D.2 Images Benchmark Pairs
[ENOT] [23] As well as for the mixtures benchmark pairs, we use the official code from
We use the same hyperparameters for this setup as the authors [23, Appendix F] except the batch size which we set to 16 (/blob/main/notebooks/Image_experiments.ipynb).
E Additional Study of Hyperparameters of Solvers
To show that the default solvers parameters described in Appendix D are already a good choice, we additionally try different values of some of the most important hyperparameters. We consider each of the solvers except [LSOT] because it is anyway known to poorly perform due to the systematic bias in its solutions [31, 32]. For the evaluation, we consider the mixtures benchmark pair with  $D = 64$  and  $\epsilon = 1$  where most of the solvers perform reasonably well. In the tables below, we use " * " to mark the hyperparameters that we use for comparisons in §4.1.
For [ENOT] solver, we consider the number of inner and outer problem iterations during the optimization and present the results in Table 10. The obtained results show that the performance increases slowly with increasing number of iterations of both types.
Inner iters
Outer iters	1	5	10	20
100	131.1	130.3	74.5	129.3
1000	28.77	47.36	25.91	20.16
10000	24.46	37.36	23.07*	18.03
Table 10: Comparison of cB $\mathbb{W}_2^2$ -UVP  $\downarrow$  (\%) for [ENOT] on mixtures benchmark pairs for  $D = 64$ ,  $\epsilon = 1$  and different hyperparameters.
For IPF-based SB solvers [MLE-SB], [DiffSB], [FB-SDE-A] and [FB-SDE-J], we try different numbers of IPF iterations and the number of samples used in each iteration. We present the results in Tables 11, 12, 13, 14. All of the IPF-based solvers learn an inversion of a diffusion process at each IPF step but they differ in the way how this is done. The typical number of IPF steps used by each algorithm is affected by this difference. The performance increases slowly with the increase of the two hyperparameters considered, at the cost of a proportional increase in iterations or in the number of samples used.
Samples per iter
IPF iters	64	128	256	512
100	23.45	24.50	16.64	14.23
1000	16.95	15.35	10.71	8.74
5000	11.55	11.24	12.96	8.41*
Table 11: Comparison of  $\mathsf{cB}\mathbb{W}_2^2$ -UVP  $\downarrow$  (\%) for  $\lfloor \mathbf{MLE} - \mathbf{SB}\rfloor$  on mixtures benchmark pairs for  $D = 64$ ,  $\epsilon = 1$  and different hyperparameters.
Samples per iter
IPF iters	64	256	512	1024
16	62.66	60.42	58.88	57.02
32	62.90	59.42	57.76*	55.08
64	62.84	59.46	57.78	55.01
Table 12: Comparison of cBW $_2$ ²-UVP  $\downarrow$  (\%) for [DiffSB] on mixtures benchmark pairs for  $D = 64$ ,  $\epsilon = 1$  and different hyperparameters.
Samples per iter
IPF iters	64	256	512
15000	173.16	163.04	160.5*
30000	168.86	165.06	156.5
Table 13: Comparison of cB $\mathbb{W}_2^2$ -UVP  $\downarrow$  (\%) for  $\lfloor \mathbf{FB - SDE - J}\rfloor$  on mixtures benchmark pairs for  $D = 64$ ,  $\epsilon = 1$  and different hyperparameters.
Samples per iter
IPF iters	64	256	512	1024
16	40.86	40.43	39.76	37.74
32	40.44	38.90	38.36*	35.46
64	40.00	38.86	38.31	35.4
Table 14: Comparison of cBW $_2$ ²-UVP  $\downarrow$  (\%) for [FB-SDE-A] on mixtures benchmark pairs for  $D = 64$ ,  $\epsilon = 1$  and different hyperparameters.
For  $\lfloor$ SCONES $\rfloor$  and  $\lfloor$ EgNOT $\rfloor$  solvers, we consider the number of Langevin steps and the Langevin step size and present the results in Table 15 and Table 16. For  $\lfloor$ SCONES $\rfloor$  the results obtained show that the performance increases slowly with increasing Langevin steps and decreasing Langevin step size. For  $\lfloor$ EgNOT $\rfloor$  the trends are slightly different, since the optimal Langevin step size seems to be in the interval [0.1, 0.2]. Anyway, our selected parameters are reasonable ones because specifying an enormously large number of Langevin steps for these solvers is sort of impractical.
Finally, for  $\lfloor \mathbf{NOT}\rfloor$  we consider the number of inner problem steps and the hidden size of the used neural network (conditional normalizing flow). We present results in Table 17.
Langevin steps
Langevin step size	64	256	512	1024
10-4	92.35	89.17	86.48	86.33*
10-3	93.51	90.41	88.22	87.74
Table 15: Comparison of cB $\mathbb{W}_2^2$ -UVP  $\downarrow$  (\%) for [SCONES] on mixtures benchmark pairs for  $D = 64$ ,  $\epsilon = 1$  and different hyperparameters.
Langevin steps
Langevin step size	100	200	500	1000
0.01	70.9	70.98	72.9	68.13
0.02	71.31	67.14	69.11	69.02
0.05	68.78	68.59	63.73*	56.84
0.1	64.52	57.45	52.35	51.9
0.2	58.22	60.08	58.93	41.31
Table 16: Comparison of cB $\mathbb{W}_2^2$ -UVP  $\downarrow$  (\%) for  $\lfloor \mathbf{EgNO}T\rfloor$  on mixtures benchmark pairs for  $D = 64, \epsilon = 1$  and different hyperparameters.
Hidden size
Inner steps	64	128	192	256	320	384	448	512
1	93.14	167.05	149.52	189.0	89.1	161.66	176.43	175.67
5	82.64	86.09	82.18	190.04	147.31	105.46	103.5	150.76
10	163.47	146.68	53.26	137.47	100.84	171.65	115.84	126.96
100	18.68	21.4	14.64	18.08	16.66	20.64*	18.71	15.15
200	61.99	52.74	58.63	53.89	52.44	55.3	55.02	54.75
Table 17: Comparison of cB $\mathbb{W}_2^2$ -UVP  $\downarrow$  (\%) for  $\lfloor \mathbf{NOT} \rceil$  on mixtures benchmark pairs for  $D = 64$ ,  $\epsilon = 1$  and different hyperparameters.
Discussion. From the results it can be seen that for the most solvers' dependence on the considered hyperparameters is almost monotonic and the hyperparameters chosen for the solver comparison on the mixtures setup are in the region where the metric growth is almost saturated.
F Qualitative Evaluation of the Drift Learned with SB methods
Our benchmark primarily aimed at quantifying the recovered conditional EOT plan  $\widehat{\pi} (\cdot |x)$ . Thanks to our Proposition 3.5, our benchmark provides not only the ground truth conditional EOT plan  $\pi^{*}(\cdot |x)$ , but the optimal SB drift  $v^{*}(x,t)$  as well. This means that for SB solvers we may additionally compare their recovered SB drift  $\widehat{v}$  with the ground truth drift  $v^{*}$ . Here we do this for [MLE-SB], [DiffSB], [ENOT], [FB-SDE-A], [FB-SDE-J] solvers by using our mixtures pairs.
METRICS. Recall that  $T_{v^*}$  is the Schrödinger bridge (10) and let  $T_{\widehat{v}}$  denote the learned process:
$$
d X _ {t} = \widehat {v} (x, t) d t + \sqrt {\epsilon} d W _ {t}, \quad X _ {0} \sim \mathbb {P} _ {0}.
$$
Both  $T_{v^*}$  and  $T_{\widehat{v}}$  are diffusion processes which start at distribution  $\mathbb{P}_0$  at  $t = 0$  and have fixed volatility  $\epsilon$ . Their respective drifts are  $v^*$  and  $\widehat{v}$ . For each time  $t \in [0,1]$ , consider
$$
\mathcal {L} _ {\mathrm {f w d}} ^ {2} [ t ] \stackrel {\text {d e f}} {=} \mathbb {E} _ {T _ {v ^ {*}}} \| v ^ {*} (X _ {t}, t) - \widehat {v} (X _ {t}, t) \| ^ {2}, \tag {31}
$$
$$
\mathcal {L} _ {\mathrm {r e v}} ^ {2} [ t ] \stackrel {\text {d e f}} {=} \mathbb {E} _ {T _ {\widehat {v}}} \| v ^ {*} (X _ {t}, t) - \widehat {v} (X _ {t}, t) \| ^ {2}. \tag {32}
$$
which are the expected squared differences between the ground truth  $v^{*}$  and learned  $\widehat{v}$  drifts at the time  $t$ . In (31), the expectation is w.r.t.  $X_{t}$  coming from the true SB trajectories of  $T_{v^{*}}$ , while in (32) - w.r.t. the learned trajectories from  $T_{\widehat{v}}$ . Reporting this metric for all the time steps, all the mixtures pairs and solvers would be an overkill. In what follows, we use this metric for quantitative analysis.
First, for  $D = 16$  and  $\epsilon \in \{0.1, 10\}$ , we plot these metrics (as a function of time  $t$ ). The results for all the solvers are shown in Figure 9. Second, we provide Table 19 where for  $D \in \{2, 16, 64, 128\}$  and  $\epsilon \in \{0.1, 10\}$  report  $\mathcal{L}^2$  metrics averaged over  $t \in [0, 1]$ . Namely, we report
$$
\operatorname {K L} \left(T _ {v ^ {*}} \| T _ {\widehat {v}}\right) \stackrel {\text {d e f}} {=} \frac {1}{2 \epsilon} \int_ {0} ^ {1} \mathcal {L} _ {\mathrm {f w d}} ^ {2} [ t ] d t \quad \text {a n d} \quad \operatorname {R K L} \left(T _ {v ^ {*}} \| T _ {\widehat {v}}\right) \stackrel {\text {d e f}} {=} \frac {1}{2 \epsilon} \int_ {0} ^ {1} \mathcal {L} _ {\mathrm {r e v}} ^ {2} [ t ] d t. \tag {33}
$$
(a)  $\frac{1}{2\epsilon}\mathcal{L}_{\mathrm{fwd}}^2 [t]$  for  $(D,\epsilon) = (16,0.1)$
(c)  $\frac{1}{2\epsilon}\mathcal{L}_{\mathrm{fwd}}^2 [t]$  for  $(D,\epsilon) = (16,1)$
Figure 9:  $\mathcal{L}^2$  metrics between the ground truth drift  $v^{*}$  and the drift  $\widehat{v}$  learned by SB solvers.
(b)  $\frac{1}{2\epsilon}\mathcal{L}_{\mathrm{rev}}^2 [t]$  for  $(D,\epsilon) = (16,0.1)$
(d)  $\frac{1}{2\epsilon}\mathcal{L}_{\mathrm{rev}}^2 [t]$  for  $(D,\epsilon) = (16,1)$
We write "KL" and "RKL" not by an accident. Thanks to the well-celebrated Girsanov's theorem, these are indeed the forward and reverse KL divergences between processes  $T_{v^*}$  and  $T_{\widehat{v}}$ .
In all the SB solvers, we consider 200 time discretization steps  $t = \left\{\frac{1}{200}, \frac{2}{200}, \dots, 1\right\}$  for their training. During testing, we evaluate  $\mathcal{L}^2$  metrics (31) and (32) on the same time steps. To estimate (31) and (32), we use  $10^5$  samples  $X_{t}$  which are taken from random trajectories of processes  $T_{v^*}$  and  $T_{\widehat{v}}$ . These trajectories are simulated via the standard Euler-Maruyama method.
ε=0.1	ε=1
D=2	D=16	D=64	D=128	D=2	D=16	D=64	D=128
[ENOT]	0.61	5.49	6.59	10.36	0.86	1.64	11.43	37.53
[DiffSB]	6.96	12.89	-	-	12.28	>1000	>1000	>1000
[FB-SDE-A]	6.9	11.08	-	-	10.59	>1000	>1000	>1000
[FB-SDE-J]	3.02	5.02	9.60	28.85	18.79	44.79	629.28	>1000
[MLE-SB]	0.62	2.63	4.76	7.86	0.96	1.86	9.66	34.95
Table 18: Forward KL between the ground truth SB process  $T_{v^*}$  and the process  $T_{\hat{v}}$  learned with SB solvers on our mixtures benchmark pairs.
ε=0.1	ε=1
D=2	D=16	D=64	D=128	D=2	D=16	D=64	D=128
[ENOT]	72.86	78.98	135.29	221.26	18.40	49.65	177.02	348.05
[DiffSB]	11.85	21.16	-	-	121.43	>1000	>1000	>1000
[FB-SDE-A]	12.29	19.40	-	-	100.22	>1000	>1000	>1000
[FB-SDE-J]	8.03	12.11	17.16	49.32	64.37	123.68	>1000	>1000
[MLE-SB]	18.03	28.24	163.34	254.16	22.80	86.07	296.97	636.27
Table 19: Reverse KL between the ground truth SB process  $T_{v^*}$  and the process  $T_{\widehat{v}}$  learned with SB solvers on our mixtures benchmark pairs.
DISCUSSION. Interestingly, we see that the forward KL divergence shows a smoother behaviour than the RKL for almost all SB solvers. According to our evaluation, the [MLE-SB] and [ENOT] solvers mostly beat every other solver in the forward KL metric. At the same time, the RKL metric of [ENOT] is surprisingly the worst. While we make all these observations, we do not know how
to explain them. We hope that the question of the interpretation of the KL and RKL values will be addressed in future SB studies.
G Potential Societal Impact
Our proposed approach deals with generative models based on Entropic Optimal Transport and Schrödinger Bridge principles. Such models form and emergent subarea in the field of machine learning research and could be used for various purposes in the industry including image manipulation, artificial content rendering, graphical design, etc. Our benchmark is a step towards improving the reliability, robustness and transparency of these models. One potential negative of our work is that improving generative models may lead to transforming some jobs in the industry.
H Building Benchmarks from Real Data
In this section, we present a simple heuristic recipe to build benchmark pairs similar to some given real-world data. To illustrate the recipe, we consider toy 2D data example and several single-cell datasets [36, 8]. Code and data for the experiments in this section can be found in the benchmark Construction/examplesdata folder of our repository.
H.1 Recipe for Building Benchmark Pairs form Data.
For constructing distribution pairs similar to some given data, we consider a pair of original and target datasets obtained from the true distributions  $\mathbb{P}_0$  and  $\mathbb{P}_1$ , respectively. We heuristically initialize the LSE potential (15)  $f^{*}(y) = \epsilon \log \sum_{n=1}^{N} w_{n} \mathcal{Q}(y|b_{n}, \epsilon^{-1}A_{n})$  with  $b_{n}$  as cluster centers obtained from the K-means clustering algorithm applied to the target data from  $\mathbb{P}_1$ . The weights  $w_{n}$  are chosen to be  $1/N$  and matrices  $A_{n} = \lambda I$  are diagonal where  $\lambda$  is a manually-chosen parameter (shared between all  $A_{n}$ ). For any  $x$  the conditional plan  $\pi^{*}(\cdot|x)$  for LSE potential  $f^{*}$  is just a Gaussian mixture and the mean of each its component is largely determined by  $b_{n}$  (Proposition 3.3). We empirically found that the resulting constructed distribution  $d\widehat{\mathbb{P}}_1(y) = d\pi_1^*(y) = \int d\pi^*(y|x)d\mathbb{P}_0(x)$  from  $\mathbb{P}_0$  resembles the Gaussian mixture approximation of the target dataset if one managed to find proper value of  $\lambda$ .
In the rest of this section, we use the described recipe to construct benchmark pairs from data to show that the LSE parameterization of the potential provides a wide class of EOT/SB solutions and even allows constructing a benchmark similar to real data.
H.2 Benchmark Pairs for 2D data.
Code and data for the experiment described in this section can be found in the folder benchmark_construction/examples/2d_data of our repository.
To begin with, we present the results of constructing a benchmark pair from 2D data. We consider a Gaussian distribution  $\mathbb{P}_0$  as the source distribution and two moons  $\mathbb{P}_1$  as the target distribution. We aim to use the previously described recipe §H.1 to find parameters of the LSE potential to construct an EOT solution between  $\mathbb{P}_0$  and an approximation of  $\mathbb{P}_1$  denoted as  $\widehat{\mathbb{P}}_1$ . Here we consider EOT with  $\epsilon = 0.05$ , use  $N = 100$  for LSE potentials, and choose  $\lambda = 50$ . The result is in Figure 10.
As seen from the figure, the constructed target benchmark distribution  $\widehat{\mathbb{P}}_1$  is similar to the target distribution  $\mathbb{P}_1$ . In turn, the EOT plan maps  $x \sim \mathbb{P}_0$  to the close regions of the target distribution.
H.3 Single-cell RNA Data
Code and data for the experiment described in this section can be found in the folder benchmark_construction/examples/single_cell_rna of our repository.
We consider the same setup as in [36, §5.2]. We use their data from the supplementary materials. The provided data displays the progression of human embryonic stem cells as they differentiate from embryoid bodies into a range of cell types, such as mesoderm, endoderm, neuroectoderm, and neural
4https://openreview.net/forum?id=d3QNWD_pcFv
Figure 10: Gaussian  $\rightarrow$  Two Moons benchmark pair.
crest, throughout a span of 27 days. The cell samples (approximately 2000 ones per each time period) were gathered at five distinct intervals ( $t_0$ : day 0 to 3,  $t_1$ : day 6 to 9,  $t_2$ : day 12 to 15,  $t_3$ : day 18 to 21,  $t_4$ : day 24 to 27). These collected cells were evaluated via scRNAseq, subjected to quality control filtering, and then projected onto a 5-dimensional feature space utilizing principal component analysis (PCA).
To construct the benchmark pair using the LSE potential, we consider  $N = 250$ ,  $\epsilon = 100$  and  $\lambda = 100$  and employ the train data at times  $t_0$  and  $t_4$ . Then we use the constructed benchmark plan  $\pi^{*}(\cdot |x)$  to map source data at time  $t_0$  to the data at time  $t_4$  and obtain benchmark target distribution samples  $\widehat{\mathbb{P}}_1$ . Finally, we fit TSNE [51] to the combined dataset of samples from  $\mathbb{P}_1$  and  $\widehat{\mathbb{P}}_1$  and then plot their projections in Figure 11. The resulting plots are very similar, confirming that the constructed benchmark target data resembles the considered single-cell target data.
Figure 11: TSNE visualization of Single-cell RNA target data and our constructed target data.
H.4 Single-cell Drugs Data
Code and data for the experiment described in this section can be found in the folder benchmark_construction/examples/single_cell_drugs of our repository.
Method	scGen	cAE	CellOT [8]	EOT Benchmark (ours)
MMD↓	0.0241	0.0074	0.0013	0.0036
Table 20: MMD $\downarrow$  distances (on the test data) between the observed perturbed cells  $\mathbb{P}_1$  and predicted responses from control cells  $\widehat{\mathbb{P}}_1$ .
In [8], the authors consider the problem of predicting single-cell drug responses for drugs with different molecular effects, using melanoma cell lines profiled by 4i technology (single-cell technology). Utilizing a blend of two melanoma tumor cell lines at a 1:1 ratio, a total of 21,650 cells were imaged. Within this dataset, 11,526 cells existed in the untreated control state, 2,364 received Erlotinib treatment, 2,650 underwent Imatinib treatment, 2,683 were subjected to Trametinib treatment, and
2,417 were treated with a combination of Trametinib and Erlotinib. After preprocessing, each cell is described by 78 features. The train-test split with each drug is 80:20.
In this example, we consider cell data before treatment  $(\mathbb{P}_0)$  and after treatment with Erlotinib  $(\mathbb{P}_1)$ . For the construction of the benchmark pair using an LSE potential, we consider  $N = 250$ ,  $\epsilon = 1$  and  $\lambda = 20$ . As with the single cell RNA data §H.3, we fit the TSNE [51] on a combined dataset of samples from  $\mathbb{P}_1$  and  $\widehat{\mathbb{P}}_1$  and then plot their projections in Figure 11. As seen from the visualizations, the TSNE projections of the real data and the mapped data are similar.
Figure 12: TSNE visualization of Single-cell Drugs target data and our constructed target data.
In addition, we quantitatively evaluate on the test data how well the constructed target distribution  $\widehat{P}_1$  matches the true data distribution  $\mathbb{P}_1$ . We employ the same MMD metric as the authors and present the results in Table 20. The data for the baselines scGen, cAE and the authors' method CellOT are taken from [8]. As one can see, our approach is even better than two of the baselines considered.