SlowGuess/NeurIPS_2023_OCR huggingface dataset:
it impossible to estimate the corresponding run-time numbers for the Sinkhorn-based group-lasso OT [11]. We have therefore excluded this algorithm here, but expect it to have larger (and possibly significantly larger) per-iteration runtime than the Sinkhorn-Knopp algorithm [11].
To make a fair comparison between the Sinkhorn-Knopp algorithm and ours, we cover four variants of the Sinkhorn-Knopp algorithm in the comparison. The implementation of Sinkhorn-Knopp in POT [16] computes the primal residual once every 10 iterations, since this operation is relatively costly compared with the iteration itself. Moreover, to increase the numerical stability of the algorithm and to improve the quality of the results with smaller weights of the entropic regularizer, the Sinkhorn-Knopp algorithm is improved by moving the computation to log scale. The decision to compute the primal residual or not, and to perform the computations in linear or logarithmic scale leads to the four distinct versions of the Sinkhorn-Knopp algorithm that we evaluate. Since our methods work directly on  $X$  we let them compute the objective value and the primal residual in every iteration.
The results show that our QDROT implementation has roughly the same per-iteration runtime as the standard Sinkhorn-Knopp algorithm. The per-iteration runtime of the Sinkhorn-Knopp algorithm in log scale is roughly 4.3-9.4x longer than QDROT and is roughly 2.5-4.5x longer than GLDROT.
This section describes additional experiments with RDROT using quadratic and group-lasso regularization on more datasets of various sizes and using different hyperparameters. The results are
Figure 6: Comparison of the per-iteration cost for OT solvers. We compare the Sinkhorn-Knopp Algorithm [11], and a log domain variant (with and without stopping criterion computations every 10th iteration), against QDROT (ours): Quadratically regularized DROT, and GLDROT (ours): Group-lasso regularized DROT.
consistent with the experiments that could be included in the main document and demonstrate significant advantages of our method compared to the state-of-the-art.
Quadratic regularization In the main manuscript, we compared our solver with an L-BFGS method applied to the dual problem. Two problem sizes were considered:  $m = 1000$ ,  $n = 1000$ , and  $m = 2000$ ,  $n = 3000$ . In Figure 2, we rerun the benchmark on 4 additional problem sizes. Notice that all results are consistent with those presented in the paper, showing that our results generalize to more settings beyond the ones presented in the paper.
Figure 7: Additional experiments of the RDROT (QDROT) and the dual L-BFGS method comparison for different quadratic regularization parameters and problem sizes. 50 datasets of an additional 4 problem sizes were simulated.
Group Lasso regularization To establish the computational advantages of RDROT applied to the Group-Lasso problem, we benchmarked it against Sinkhorn-based solver on simulated datasets of size  $m = 1000$ ,  $n = 1500$ , for a selection of hyperparameters. It Table 4,5,6, we rerun the benchmark on more problem sizes and hyperparameters, to test if the results generalize. Indeed, GLDROT consistently outperforms the state-of-the-art in all considered problems.
Table 4: GL experiments with  $n = {500}$  and  $m = {500}$
Method	Reg.	Runtime (s) ↓	Agg. W2 dist. ↓
Ent	GL	Median	q10	q90	Median	q10	q90
GLSK	1e-3	1e-6	1.26	1.24	2.98	0.35	0.084	4.96
1e-3	5e-4	1.23	1.22	2.93	0.351	0.0841	4.96
1e-3	5e-2	1.35	1.23	3.57	0.377	0.092	4.96
1e-2	1e-6	1.21	1.09	2.63	1.22	0.668	5.48
1e-2	5e-4	1.18	1.07	2.61	1.22	0.668	5.48
1e-2	5e-2	1.2	1.08	2.63	1.22	0.668	5.57
1e-1	1e-6	1.04	1.04	1.06	7.83	3.27	30.6
1e-1	5e-4	1.05	1.05	1.07	7.83	3.27	30.6
1e-1	5e-2	1.05	1.04	1.06	7.83	3.27	30.6
1.	1e-6	1.03	0.517	1.04	45.3	9.74	212
1.	5e-4	1.02	0.512	1.03	45.3	9.74	212
1.	5e-2	1.03	0.514	1.03	45.3	9.74	212
GLDROT	1e-6	0.0924	0.0519	0.16	0.0618	0.028	0.309
5e-4	0.0553	0.0354	0.106	0.0801	0.0357	0.319
5e-2	0.0127	0.00889	0.0351	0.672	0.321	2.87
Table 5:  $m = {1000},n = {1000}$
Method	Reg.	Runtime (s) ↓	Agg. W2 dist. ↓
Ent	GL	Median	q10	q90	Median	q10	q90
GLSK	1e-3	1e-6	4.97	4.91	9.27	1.66	0.405	11
1e-3	5e-4	4.9	4.82	9.19	1.66	0.405	11
1e-3	5e-2	5.52	4.91	15.9	1.64	0.389	11
1e-2	1e-6	4.9	4.79	5.53	3.06	0.926	8.97
1e-2	5e-4	4.79	4.66	5.5	3.06	0.926	8.97
1e-2	5e-2	4.88	4.73	5.55	3.15	0.926	9.03
1e-1	1e-6	4.71	4.68	4.73	20.6	8.97	74.3
1e-1	5e-4	4.74	4.7	4.78	20.6	8.97	74.3
1e-1	5e-2	4.76	4.73	4.82	20.6	8.97	74.3
1.	1e-6	4.61	4.57	4.64	130	37.5	403
1.	5e-4	4.62	4.6	4.67	130	37.5	403
1.	5e-2	4.74	4.71	4.77	130	37.5	403
GLDROT	1e-6	0.109	0.0735	0.146	0.0999	0.0424	0.304
5e-4	0.0598	0.045	0.0763	0.138	0.0632	0.44
5e-2	0.0281	0.0232	0.0374	3.16	1.14	9.01
Table 6:  $m = {1000},n = {1500}$
Method	Reg.	Runtime (s) ↓	Agg. W2 dist. ↓
Ent	GL	Median	q10	q90	Median	q10	q90
GLSK	1e-3	1e-6	3.82	3.75	9.03	0.311	0.0657	6.48
1e-3	5e-4	3.8	3.75	9	0.311	0.0657	6.48
1e-3	5e-2	3.8	3.77	10.6	0.345	0.0665	6.48
1e-2	1e-6	3.78	3.73	6.8	1.21	0.69	5.47
1e-2	5e-4	3.8	3.76	6.82	1.21	0.69	5.47
1e-2	5e-2	3.82	3.79	6.88	1.21	0.69	5.47
1e-1	1e-6	3.7	3.67	3.72	8.24	3.47	31.6
1e-1	5e-4	3.75	3.73	3.78	8.24	3.47	31.6
1e-1	5e-2	3.77	3.73	3.81	8.24	3.47	31.6
1.	1e-6	3.73	1.86	3.77	45.5	9.92	218
1.	5e-4	3.69	1.85	3.73	45.5	9.92	218
1.	5e-2	3.75	1.88	3.78	45.5	9.92	218
GLDROT	1e-6	0.113	0.0758	0.147	0.0475	0.0215	0.137
5e-4	0.0745	0.0549	0.0951	0.0529	0.0245	0.163
5e-2	0.0232	0.0178	0.0288	0.331	0.154	1.38
G Model specification of the generative Adversarial Model
We adopted similar network structures and used the same loss function as in [35] and performed several experiments with image generation based on the MNIST and CIFAR10 datasets.
To define the loss, we need the following sample-based OT-cost.
$$
\begin{array}{l} \mathcal {W} _ {c, h} (\mathbf {X}, \mathbf {Y}) = \mathrm {O T} _ {h} (C _ {\mathbf {X}, \mathbf {Y}}) \\ := \inf  _ {M \in \mathbb {R} _ {+} ^ {m \times n}} \left\{\langle C _ {\mathbf {X}, \mathbf {Y}}, M \rangle + h (M): M \mathbf {1} _ {n} = m ^ {- 1} \mathbf {1} _ {m}, M ^ {\top} \mathbf {1} _ {m} = n ^ {- 1} \mathbf {1} _ {n} \right\}. \\ \end{array}
$$
Here  $m$  and  $n$  are the batch sizes corresponding to  $\mathbf{X}$  and  $\mathbf{Y}$ , and  $C_{\mathbf{X},\mathbf{Y}}$  is a matrix with pairwise distances between the samples of  $\mathbf{X}$  and  $\mathbf{Y}$ . The mini-batch energy distance [35] is then given by
$$
\begin{array}{l} \mathcal {L} _ {h} = \mathcal {W} _ {c, h} (\mathbf {X}, \mathbf {Y}) + \mathcal {W} _ {c, h} \left(\mathbf {X} ^ {\prime}, \mathbf {Y}\right) + \mathcal {W} _ {c, h} \left(\mathbf {X}, \mathbf {Y} ^ {\prime}\right) + \mathcal {W} _ {c, h} \left(\mathbf {X} ^ {\prime}, \mathbf {Y} ^ {\prime}\right) \\ - 2 \mathcal {W} _ {c, h} (\mathbf {X}, \mathbf {X} ^ {\prime}) - 2 \mathcal {W} _ {c, h} (\mathbf {Y}, \mathbf {Y} ^ {\prime}). \\ \end{array}
$$
Here,  $\mathbf{X}$ ,  $\mathbf{X}'$  are two independent mini-batches from real data, while  $\mathbf{Y}$ ,  $\mathbf{Y}'$  are two independent mini-batches generated from the generator.
In the experiments, we used the cosine similarity to parameterize the cost, and we replaced the OT solver used in [35] - the Sinkhorn-Knopp algorithm, by the PyTorch wrapper for our OT solver with the quadratic regularizer.
The model structures are adapted from the ones in [35]. Weight normalization is used to construct the parameters of the models [34], and the activation functions are gated linear units [13] and concatenated ReLUs [38]. We train the model with the Adam optimizer [22] using an initial learning rate of  $3 \times 10^{-4}$ , with  $\beta_{1} = 0.5$  and  $\beta_{2} = 0.999$ . The batch size is 1024 for the MNIST experiment and 2048 for the CIFAR10 experiment. We update the generator three times for every discriminator update. As for the parameters of our OT solver, we use  $\epsilon = 10^{-4}$  as the stopping criterion and  $\lambda = 10^{-3}$  as the weight of the quadratic regularizer. The results for MNIST and CIFAR10 are detailed below.
Different combinations of hyperparameters, model structures and schedules for the weight of the regularizers and the stopping criterion (the OT solver can allowed to produce low-accuracy solutions at the beginning of the training, for example) may potentially improve the results of the generated samples, but since GANs are not the focus of this paper, we leave such refinements as future work.
G.1 MNIST Architecture and Results
Operation	Activation	Kernel Size	Stride	Padding	Output Shape
Sample z	[4]
Linear	GLU	[128·7·7]
Reshape	[128, 7, 7]
Upsample × 2	[128, 14, 14]
2D Convolution	GLU	[5, 5]	1	same	[128, 14, 14]
Upsample × 2	[128, 28, 28]
2D Convolution	GLU	[5, 5]	1	same	[64, 28, 28]
2D Convolution	tanh	[5, 5]	1	same	[1, 28, 28]
Table 7: Generator Architecture for MNIST
Opertion	Activation	Kernel Size	Stride	Padding	Output Shape
2D Convolution	CReLU	[5, 5]	1	same	[128, 28, 28]
2D Convolution	CReLU	[5, 5]	1	same	[128, 28, 28]
2D Convolution	CReLU	[5, 5]	2	2	[256, 14, 14]
2D Convolution	CReLU	[5, 5]	2	2	[256, 7, 7]
Flatten	[256 · 7 · 7]
L2 Normalization	[256 · 7 · 7]
Table 8: Discriminator Architecture for MNIST
Figure 8: Generated samples at 100th epoch for MNIST (Original size for each sample is  $28 \times 28$  px)
The generated samples in Figure 8 are clear and shape, and no mode collapse can be observed (digits from 0 to 9 with various variants can be found in the figure). Compared with the results in [17] (Figure 3), our results show higher quality.
G.2 CIFAR10 Architecture and Results
Opertion	Activation	Kernel Size	Stride	Padding	Output Shape
Sample z	[100]
Linear	GLU	[1024·4·4]
Reshape	[1024, 4, 4]
Upsample × 2	[1024, 8, 8]
2D Convolution	GLU	[5, 5]	1	same	[512, 8, 8]
Upsample × 2	[512, 16, 16]
2D Convolution	GLU	[5, 5]	1	same	[256, 16, 16]
Upsample × 2	[256, 32, 32]
2D Convolution	GLU	[5, 5]	1	same	[128, 32, 32]
2D Convolution	tanh	[5, 5]	1	same	[3, 32, 32]
Table 9: Generator Architecture for CIFAR10
Opertion	Activation	Kernel Size	Stride	Padding	Output Shape
2D Convolution	CReLU	[5, 5]	1	same	[256, 32, 32]
2D Convolution	CReLU	[5, 5]	2	2	[512, 16, 16]
2D Convolution	CReLU	[5, 5]	2	2	[1024, 8, 8]
2D Convolution	CReLU	[5, 5]	2	2	[2048, 4, 4]
Flatten	[2048 · 4 · 4]
L2 Normalization	[2048 · 4 · 4]
Table 10: Discriminator Architecture for CIFAR10
Figure 9: Generated samples at 1900th epoch for CIFAR10 (Original size for each sample is  $32 \times 32$  px)
In Figure 9, most samples are recognizable, and the samples cover the 10 categories of the CIFAR10 dataset with enough diversity. Our generated samples have similar quality as the ones reported in [35] (Figure 4) which uses the Sinkhorn-Knopp algorithm as the OT solver in the mini-batch energy loss.