Lightspeed Geometric Dataset Distance via Sliced Optimal Transport

We would like to thank the reviewer for the time and constructive feedback.

Q11. Please provide large-scale experiments where OTDD is too expensive, and s-OTDD excels.

A11. In the paper, we showed that OTDD cannot be used when the dataset size reaches more than 30,000 samples on MNIST and CIFAR10 in Figure 2. In addition, we added new experiments on the large-scale Tiny-ImageNet dataset, where the OTDD baseline is not feasible to compute in A14.

Q12. The datasets used in the work to compute the distance tend to be very simplistic, such as NIST-type datasets.

A12. We did conduct text classification with various datasets such as AG News, DBPedia, Yelp Reviews, Amazon Reviews, and Yahoo Answers (Section 4.2). In addition, we conducted experiment on Tiny ImageNet and CIFAR-10 (32×32 resolution). We also refer the reviewer to A14 where we added a new experiment on Tiny-ImageNet of size 224×224.

Q13. The reason to use Pearson Correlation between OTDD and s-OTDD for evaluation is not clear.

A13. s-OTDD can be seen as an alternative solution for OTDD when dealing with a large scale datasets. Therefore, measuring correlation with OTDD strengthens the claim that s-OTDD can replace OTDD while being more efficient.

Q14. A large-scale evaluation on high dimensional datasets is required to demonstrate the effectiveness of the method over OTDD.

A14. We have added a large-scale transfer learning experiment in Tiny-ImageNet dataset (224x224 resolution) where computing OTDD is infeasible. To compute dataset distances, we sampled 5,000 examples from each sub-dataset. The result of the experiment is visible in Figure 16 at this link https://imgur.com/a/hHvg2T2. In the figure, we see that s-OTDD has a relatively good correlation with the classification accuracy.

Q15. Why is OTDD a natural choice for dataset distance? Why should a new dataset distance try to be correlated with OTDD? Why should the correlation be the Pearson correlation (wouldn't a rank correlation ...)?

A15. OTDD is a good dataset distance since it is model-agnostic, does not involve training, can compare datasets even if their label sets are completely disjoint. Similar to OTDD, s-OTDD is also model-agnostic (requires no modeling on data), does not need to estimate any parameters, and can handle label sets that are completely disjoint. Since s-OTDD is proposed as an alternative solution to OTDD, it needs to have a good correlation with OTDD.

From the suggestion of the reviewer, we have added the new figures with both Pearson correlation (denoted as $r$) and Spearman's rank correlation (denoted as $\rho$) in Figure 12-14 at this link https://imgur.com/a/hHvg2T2. We observe that the values of Spearman's rank correlation and Pearson correlation are almost similar in majority of cases for s-OTDD.

Q16. The number of random projections ... What are the implications of it for larger datasets? The time complexity for the method increases to $O(n^2 \log n)$ and space complexity $O(n(n + d))$... how many projections should one be looking at for the best results?

A16. In practice, $L$ should be large enough with respect to the number of dimensions $d$ not the dataset size $n$. In particular, authors in [4] empirically show that we need $L\approx 1.22 \sqrt{d}$. In Figure 2, we showed that the computation gap between $L=1000$ and $L=10000$ is not large. Moreover, the computations for $L$ projections are independent, hence, we can utilize parallel computing to compute s-OTDD. Therefore, having a large $L$ is not a problem. In addition, we can reduce $L$ by using more advanced sampling techniques [5].

[4] Sliced Wasserstein Autoencoders, Kolouri et al. [5] Quasi-Monte Carlo for 3D Sliced Wasserstein, Nguyen et al.

Q17. ... s-OTDD can work with more data than OTDD, why is it not more predictive ... ?

A17.. In the paper, we actually compare s-OTDD and OTDD with the same number of data samples. We want to convey that s-OTDD is much faster than OTDD while being comparable in performance. In Figure 2 and in A14, we showed that OTDD cannot be used when having large or high-dimensional datasets.

Q18. Is s-OTDD predictive of metrics other than classification accuracy?

A18. We can compute the correlation between s-OTDD and any function of two datasets. We added a new experiments to show that s-OTDD also correlates well with other performance metrics such as Precision, Recall, and F1 Score in Figure 15 https://imgur.com/a/hHvg2T2.

Q19. What types of datasets/tasks (...) can s-OTDD be used to measure distances for?

A19. s-OTDD can also be interpreted as a distance between distributions over distributions. Therefore, s-OTDD can also be adapted to compare distributions over point clouds, distributions over histograms, distributions over 3D shapes, multi-label datasets, and so on.

We appreciate the time and constructive feedback of the reviewer. We would like to extend the discussion with the reviewer as follows:

Q6. Proposition 1 gives two conditions for the MTP to be injective, and the metric properties of s-OTDD are based on the injectivity of MTP. Does the MTP used in the experiments satisfy the conditions? Please justify the numerical implementation of MTP, including the assumption of infinite number of moments, the use of $\sigma(\Lambda^k)$ and how the metric properties are affected.

A6. Checking existence of all moments of a high-dimensional distribution is impractical due to a high computational cost. Moreover, we only observe random samples from the unknown distribution which makes the estimation of moments are also random. Therefore, in experiments, we implicitly assume that the underlying distributions of datasets have infinite moments. This assumption is motivated by the fact that one-dimensional projection becomes a Gaussian distribution in high-dimension [1,2] and Gaussian distribution has infinite number of moments. We then use zero truncated Poisson distribution (as reported at line 309-310) for the value of moments since it is a distribution over infinite countable natural numbers. If the underlying distributions of datasets do have infinite moments, MTP is injective and s-OTDD is a metric. We would like to recall that s-OTDD is still a pseudo distance without injectivity of the MTP i.e., s-OTDD satisfies the triangle inequality, symmetry, non-negativity, and one direction of identity (is 0 when two datasets are the same). Therefore, s-OTDD is still a meaningful discrepancy for datasets as demonstrated through our experiments.

[1] Fast approximation of the sliced-Wasserstein distance using concentration of random projections, NeurIPS 2021, Nadjahi et al.

[2] Asymptotics of Graphical Projection Pursuit, The Annals of Statistics, 1984 Diaconis et al.

Q7. In Definition 3, do feature projection and MTP have to share the same $\psi^{(1)}$?

A7. It is actually a typo in our Eq. (9). The summation in the second term should involve $\psi^{(i+1)}$. We have fixed this typo in the revision. In particular, we have the following revised definition of the data point projection

$\mathcal{DP}^k_{\psi,\theta,\lambda,\phi}(x,q_y) = \psi^{(1)} \mathcal{FP}_\theta(x)$

$+ \sum_{i=1}^k\psi^{(i+1)} \mathcal{MTP}_{\lambda^{(i)},\phi}(q_y)$,

where $\psi=(\psi^{(1)},\psi^{(2)},\ldots,\psi^{(k+1)}) \in \mathbb{S}^k, \theta \in \Theta, \lambda =(\lambda^{(1)},\ldots,\lambda^{(k)}) \in \Lambda^{k}, \phi \in \Phi$.

Q8. In Corollary 1, if both feature projection and MTP are injective, why the data point projection, as the sum of them, is also injective? Could the authors explain why the data point projection is a composition of Radon transform projection and the MTP?

A8. Thank you for your questions. From the revised definition in A7. let $\mathcal{FP}_\theta(x)=t_1$

and $\mathcal{MTP}_{\lambda^{(i)},\phi}(q_y)=t_i$ for $i=2,\ldots,k$, we can rewrite the data point projection as:

$\mathcal{DP}^k_{\psi,\theta,\lambda,\phi}(x,q_y) = \psi^{(1)} t_1+ \sum_{i=2}^k\psi^{(k)} t_i = \psi^\top t = \mathcal{R}_\psi(t)$,

where $\psi=(\psi^{(1)},\ldots,\psi^{(k)})$, $t=(t_1,\ldots,t_k)$, and $\mathcal{R}_\psi(t)$ is the Radon Transform of $t$ with projection parameter $\psi$ as defined at line 162. Therefore, the data point projection is nothing but the Radon Transform of the stacked output of the feature projections and moment transform projections. Since the composition of injective functions is also injective, the data point projection is injective given the injectivity of the feature projection and the moment transform projections.

Q9. Can you explain more about the feature projection choices for each of the dataset?

A9. As written at line 194, we choose the feature projection based on the prior knowledge about the feature space. For example, if we believe that the feature space is an Euclidean space, we can use the Radon projection (linear projection). If we work with a manifold, we can use geodesic projection [3]. If we work with images, we can use the convolution projection [4]. In our experiments, for the NIST dataset and the text classification dataset, we apply linear projection by default. For image data, we use convolutional projection. We conducted a new comparison between linear and convolutional projections, available at https://imgur.com/a/1NVv5AC. The results show that convolution-based projections not only require fewer projections but also tend to exhibit a stronger positive correlation.

[3] Sliced-Wasserstein Distances and Flows on Cartan-Hadamard Manifolds, JMLR, 2025, Bonet el al.

[4] Revisiting Sliced Wasserstein on Images: From Vectorization to Convolution, NeurIPS, 2022, Nguyen et al.

Q10. Typos

A10. We have fixed them in the revision.

First, we would like to thank the reviewer for the time and the feedback. We extend the discussion as follows:

Q1. The line below Eq. (8), empirical distribution is missing the variable.

A1. Thank you for pointing out the typo. We have fixed the typo in the revision.

Q2. In Eq. (9), is the first $\psi^{(1)}$ in the first term, the same as the $\psi^{(k)}$ with $k = 1$ in the second term?

A2. It is actually a typo in our Eq. (9). The summation in the second term should involve $\psi^{(i+1)}$. We have fixed the typo in the revision. In particular, we have the following revised definition of the data point projection:

$\mathcal{DP}^k_{\psi,\theta,\lambda,\phi}(x,q_y) = \psi^{(1)} \mathcal{FP}_\theta(x)$

$+\sum_{i=1}^k \psi^{(i+1)} \mathcal{MTP}_{\lambda^{(i)},\phi}(q_y)$

where $\psi=(\psi^{(1)},\psi^{(2)},\ldots,\psi^{(k+1)}) \in \mathbb{S}^k, \theta \in \Theta, \lambda =(\lambda^{(1)},\ldots,\lambda^{(k)}) \in \Lambda^{k}, \phi \in \Phi$.

We would like to elaborate more on Corollary 1 with the revised definition. Let $\mathcal{FP}_\theta(x)=t_1$

and $\mathcal{MTP}_{\lambda^{(i)},\phi}(q_y)=t_i$ for $i=2,\ldots,k$, we can rewrite the data point projection as:

$\mathcal{DP}^k_{\psi,\theta,\lambda,\phi}(x,q_y) = \psi^{(1)} t_1+ \sum_{i=2}^k\psi^{(k)} t_i = \psi^\top t = \mathcal{R}_\psi(t)$,

where $\psi=(\psi^{(1)},\ldots,\psi^{(k)})$, $t=(t_1,\ldots,t_k)$, and $\mathcal{R}_\psi(t)$ is the Radon Transform of $t$ with projection parameter $\psi$ as defined at line 162. Therefore, the data point projection is simply the Radon Transform of the stacked output of the feature projections and moment transform projections. Since the composition of injective functions is also injective, the data point projection is injective given the injectivity of the feature projection and the moment transform projections.

First, we would like to thank the reviewer for the review and feedback. We would like to answer questions from the reviewer as follows:

Q3. A weakness of the proposed sliced OTDD is the extra hyper-parameter - the number of moment $\lambda$ as in (7). In the code, it’s set to 5. I didn’t find it discussed in the paper except in [246] where authors only briefly mentioned the principle of choosing $\lambda$. Its impact on the projection and hence the distances is unknown as well.

A3. Thank you for your insightful comments. We would like to recall that the number of moments is denoted as $k$ and we did report the choice of $k=5$ at line 307-308 in the paper. For the parameter $\lambda$, it is the order of a moment. We would like to extend the discussion on both $k$ and $\lambda$ as follows:

For $k$, it plays the role of selecting the number of moments for the label projection. In particular, the higher value of $k$, the more moments information of the label distribution is gathered into the data point projection. For $\lambda$, it plays the role of choosing information of distributions to extract. For example, $\lambda=1$ leads to the information of centerness while $\lambda=2$ leads to the information of spread. In the paper, $\lambda_i$ for $i=1,\ldots,k$ are random which followed the zero-truncated Poisson distribution with the rate parameter equals $i$ i.e., the mean $\mathbb{E}[\lambda_i]= (ie^i)/(e^i-1) \approx i$. It means that we try to capture (in high probability) the first $k$ moments of the label distribution. In the case, where we know the true number of moments of the label distribution $\lambda^\star$, we can set $k=\lambda^*$ to capture all information of the label distribution. Nevertheless, checking existence of moments is expensive since we only access to samples from the dataset generating distribution and the size of datasets is also large. Therefore, $k$ becomes a hyperparameter. However, choosing $k$ is not a hard problem since we know that we want to have as big $k$ as possible. In practice, we can start with a big $k$ and check if the value of s-OTDD exists (not overflow). If s-OTDD does not exist, you can reduce the value of $k$ via a binary search rule to have a choice of $k$. To avoid such searching algorithm, we recommend to choose $k=5$ due to the concentration of random projection. In particular, we know that one-dimensional projection becomes a Gaussian distribution in high-dimension [1,2]. Therefore, two moments are enough to capture the information of a Gaussian distribution. Since we have finite dimension, using extra 3 moments (5 in total) must be enough.

To verify the above hypothesis, we have conducted additional ablation studies in *NIST Adaptation where we vary $k\in \\{1,2,3,4,6\\}$, resulting visually at https://imgur.com/a/NCcqYgo. We see that increasing $k$ leads to better correlation to the performance gap. Nevertheless, we found that when passing $k=3$, the correlation does not increase fast as when $k<3$. Also, we found that when setting $k$ to be too large, s-OTDD is not computable due to numerical issue i.e., overflow which might be due to the non-existence of the higher moments. We will add this discussion and the additional experiments to the revision.

[1] Fast approximation of the sliced-Wasserstein distance using concentration of random projections, NeurIPS 2021, Nadjahi et al.

[2] Asymptotics of Graphical Projection Pursuit, The Annals of Statistics, 1984 Diaconis et al.

Q4.. The $\lambda$-th scaled moment as defined in (7) is not how I usually "scale" a moment. What’s the rationale behind a constant scaler? Please discuss the impact of $\lambda$.

A4. The reason we scale $\lambda$-th moment by $\lambda!$ is to make sure the output of the data point projection is not biased toward high order moments. In particular, in the definition of the data point projection (Definition 3), the data point projection is a weighted average of the feature projection and multiple moment transform projections with different values of $\lambda$. If we do not scale the moment, the value of the data point projection will be entirely based on the value of the moment transform projection with highest value of $\lambda$. However, we know that each moment captures different information of the distribution e.g., the first moment captures the information about center, the second moment captures the information about the spread. Therefore, scaling the moment value makes the contribution of all moments "approximately" equal. The reason we say that they are "approximately" equal is that we actually imply the priority for small order moments since the factorial normalizing constant goes faster in limit than the exponential function. We will include this discussion in the revision of the paper.

Q5. provide more details for A.1.

A5. Thank you for your comments. We will expand the proofs by adding intermediate steps.