SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We finally evaluate min-SWGG in an other context: computing distances between datasets. Let  $\mathcal{D}_1 = \{(x_i^1, y_i^1)\}_{i=1}^n$  and  $\mathcal{D}_2 = \{(x_i^2, y_i^2)\}_{i=1}^n$  be source and target data sets such that  $x_i^1, x_i^2 \in \mathbb{R}^d$  are samples and  $y_i^1, y_i^2$  are labels  $\forall 1 \leq i \leq n$ . In [3], the authors compare those data sets using the Wasserstein distance with the entries of the cost matrix defined as:
$$
C _ {i j} = \left(\| \boldsymbol {x} _ {i} ^ {1} - \boldsymbol {x} _ {j} ^ {2} \| _ {2} ^ {2} + W _ {2} ^ {2} \left(\alpha_ {\boldsymbol {y} _ {i} ^ {1}}, \alpha_ {\boldsymbol {y} _ {j} ^ {2}}\right)\right) ^ {1 / 2} \tag {87}
$$
and the corresponding distance as:
$$
O T D D \left(\mathcal {D} _ {1}, \mathcal {D} _ {2}\right) = \min  _ {P \in U} \langle C, P \rangle \tag {88}
$$
where  $\alpha_{\pmb{y}}$  is the distribution of all samples with label  $\pmb{y}$ , namely  $\{\pmb{x} \in \mathbb{R}^d | (\pmb{x}, \pmb{y}) \in \mathcal{D}\}$  for  $\mathcal{D}$  being either  $\mathcal{D}_1$  or  $\mathcal{D}_2$  and  $U$  is the Birkhoff polytope which encodes the marginal constraints. Notice that cost in Eq. (87) encompasses the ground distance and a label-to-label distance. This distance is appealing in transfer learning application since it is model-agnostic. However, it can be cumbersome to compute in practice since it lays down on solving multiple OT problems (to compute the cost matrix and the OTDD). To circumvent that, [3] proposed several methods to compute the cost matrix in Eq. (87). They used the Sinkhorn algorithm (in  $\mathcal{O}(n^2)$ ) or they assumed  $\alpha_{\pmb{y}} \sim \mathcal{N}(m_{\pmb{y}}, \Sigma_{\pmb{y}})$  in order to get the WD through the Bures metric (that provides a closed form of OT for Gaussian distributions in  $\mathcal{O}(d^3)$ ), which is still prohibitive for high dimension. We challenge min-SWGG in this context.
In this experiment, we compare the following datasets: MNIST [39], EMNIST [18], FashionMNIST [67], KMNIST [17] and USPS [32]. We rely on the code of OTDD provided at https://github.com/microsoft/otdd. In order to make it compliant with the min-SWGG hypothesis, we require the empirical distributions  $\alpha_{\pmb{y}}$  to have the same number of atoms.
Fig. 19 provides results for a batch size of  $n = 40000$  samples using the Sinkhorn divergence (with a regularisation parameter of  $1e^{-1}$ ) and for min-SWGG (optimized) on batch of size 40000. We report results for a learning rate of  $1e^{-5}$ , 20 iterations and  $s$  and  $\epsilon$  to be 1 and 0.
Figure 19: OTDD results  $(\times 10^{2})$  distances for min-SWGG (left) and Sinkhorn divergence (right) for various datasets.
We check that the orders of magnitude are preserved with min-SWGG. For example OTDD(MNIST,USPS) is smaller than OTDD(MNIST,FashionMNIST) for either Sinkhorn divergence or min-SWGG as distance between labels, this validate that min-SWGG is a meaningful distance in this case scenario. Moreover in our setup, the computation cost is more expensive for Sinkhorn than for min-SWGG and totally untractable for  $W_2^2$ . On smaller batches (see Fig 20), the same observation can be made: min-SWGG is comparable (in term of magnitude) with  $W_2^2$ , Sinkhorn and the Bures approximation.
We give additional results in Fig. 20 for batches of size of  $n = 2000$  samples obtained with  $W_2^2$ , Sinkhorn divergence (setting the entropic regularization parameter to  $1e^{-1}$ ), the Bures approximation, min-SWGG (random search with  $L = 1000$  projections) and min-SWGG (optimized, with a learning rate of  $5e^{-1}$ , 50 iterations,  $s = 20$  and  $\epsilon = 0.5$ ).
Figure 20: OTDD distance with  $W_{2}^{2}$  (left), Sinkhorn (left-mid), Bures (middle) and random search min-SWGG (right-mid) and min-SWGG-optimization (right) distances between labels distribution  $\times 10^{2}$
Note that the Figs. 19 and 20 are not symmetric (OTDD(KMNIST,FashionMNIST)  $\neq$  OTDD(FashinMNIST,KMNIST)) because of the random aspect of batches.