Joint Metric Space Embedding by Unbalanced Optimal Transport with Gromov–Wasserstein Marginal Penalization

Many thanks for your positive report and the constructive comments. We revised our manuscript accordingly, especially with respect to the limiting cases and additional references. Please find our detailed answers below.

Limiting Cases: Thanks for pointing this out, which extends the limiting case ($\lambda\to\infty$) for existing isometric embeddings discussed in Prop 3.4. As you suggested, the limit ($\lambda\to0$) actually yields a fixed support barycenter, and the limit ($\lambda\to\infty$ without isometries) gives the best possible projection for both spaces separately. In the revised version, we extend Section 3 by adding a new proposition. The rigorous proof relies on a subsequence argument and the weak continuity of the EW functional similar to the proof of Prop 3.4. The proof is added to the supplement.

Extendability: From a theoretical point of view, the proposed variational formulation can be adapted to handle multiple input spaces. The main issue hereby is how to compare the different projections on $Z$. Intuitively, we may rely on pairwise Wasserstein distances; but this approach does not yield a 'tree-like' cost function for the multi-marginal reformulation. As a consequence, the employed multi-marginal Sinkhorn method (Beier et al, 2022, doi:10.1007/s10851-022-01126-7) cannot be applied. An alternative would be to utilize only the Wasserstein distance between the $k$th and $(k+1)$st projection. Since the employed multi-marginal Sinkhorn method linearly scales in the number of spaces, we expect that the resulting method remains numerically tractable for small numbers of spaces. The actual choice of the model and the numerical implementation of a multiple space variant, however, requires further studies. Thanks for pointing out this imported extension, which is discussed in the conclusion of the revised manuscript.

Multi-Marginal Reformulation: In the final version, the rewriting of the original problem is extended to give a better intuition behind the equivalence. The detailed proof of Prop 3.5 is added to the supplement.

Choice of Geometry You are right: the geometry of $Z$ has to be chosen with respect to the problem at hand. To study the sensitivity with respect to the discretization of $Z$, we repeated the experiments in Section 6.1 (unrolling, 2d embedding) with different grid sizes. The results remain stable meaning that we obtain the reported embeddings in different resolutions ($Z$-Experiment). The chosen grid, however, has a major impact on the computation time. These experiments are added to the supplement. Please notice that, besides the alignments on the sphere and torus, also the alignment of the Gaussian mixture models rely on a non-Euclidean geometry since the alignment is calculated in the space of Gaussians equipped with the Wasserstein distance itself. We agree that these experiments are only proof-of-concepts that can only point in the direction of possible applications and that real-world applications require further study.

Parameter and Regularization: To study the impact of $\lambda$ and the entropic regularization, we extend the experiments in section 6.1. As you pointed out, small $\lambda$ leads to fixed support barycenters and large $\lambda$ to GW-minimizing projections of the input spaces ($\lambda$-Experiment), which nicely relates to the theoretical limiting cases. To study the impact of the regularization, we again adapt the experiments in Section 6.1 ($\varepsilon$-Experiment). The results are added to the supplement, and we discuss the impact of entropic regularization.

Code and Reproducibility: Thanks for pointing out that the parameter ranges for the grid searches are missing. We use the same distance graph and the same entropic parameter $\varepsilon=0.001$ for all algorithms and four hyperparameter configurations per method. We test $\lambda=0.01,0.1,10,100$ for JMDS and our method. For UnionCOM, we use combinations of $\beta=1,10$ and a perplexity of 30 and 100 for the t-SNE dimensionality reduction. SCOT only relies on $\varepsilon$ and the input distances and has no additional hyperparameters. In the revised version, the ranges are added. Also, our implementations will be publicly available via Github to reproduce or reuse our implementations. Due to the anonymity requirement and a short e-mail exchange on the official guidelines, we cannot link the code here, but the link will be provided in the final version.

References: Thanks for pointing out the related references which we incorporate in the revised version. Regarding the GW dimension reduction, we also refer to (Clark, Needham, Weighill: arXiv:2405.15959v2), where the relation between GW and MDS is discussed.

We would like to thank you again for your suggestions that significantly improved the manuscript and hope that you can accept the paper now.

Many thanks for the positive evaluation and your constructive feedback! We have revised our manuscript according to your suggestions in particular with respect to quantitative experiments. Please find our detailed answers to your questions below.

Quantitative FAUST: We agree that the results shown in Figure 4 and 5 are qualitative and can only give fleeting impressions of the algorithms. Since our method and JMDS can both be derived from the same functional, we compare the achieved objectives, the Wasserstein distances between embeddings, and the GW distances between embeddings and input spaces. The experiment is repeated for each pairwise combination of the first ten shapes from the FAUST dataset and different $\lambda$. For comparison, the shapes are downsampled to 35 vertices, and both methods are regularized with $\varepsilon=0.001$. For our method, we set $Z$ according to a uniform 30$\times$30 grid on $[0, 1.3]$. The average and standard deviation of the GW distance (input space to embedding) and Wasserstein distance (between embeddings) are as follows: Quantitative FAUST. In total, we here achieve better alignments than JMDS both in terms of GW and Wasserstein. While our method requires 4-5 seconds per embedding, JMDS requires 2-3 seconds. These results are added to the revised version.

Separability: We agree that the classes in Figure 6 are not clearly separated. Nevertheless, the constructed 2d embeddings increases the KNN accuracy compared with the other methods. The message behind Figure 6 is that the embeddings of the two datasets from completely different feature spaces looks quite similar and that the classes are mapped to common regions, i.e., the embeddings have similar color codings. To improve the visual embeddings, we now equip the feature spaces with distances from the unsupervised SCOT distance estimation routine (Demetci, et al, doi:10.1089/cmb.2021.0446). The updated distances lead to the following new version of Figure 6.

Experiments: Please note that there are additional experiments in the supplement (comparison between GW, $\textrm{EW}_\lambda$ and Wasserstein in Appendix B; domain adaption experiment in Appendix C). To further improve the numerical section, we add the above quantitative study of the FAUST dataset. Moreover, we add illustrative supplementary sensitivity studies with respect to the relaxation parameter ($\lambda$-Experiment), the entropic parameter ($\varepsilon$-Experiment) and the discretization of the joint space ($Z$-Experiment).

Minor Comments:

• Please notice that the push-forward is introduced in the first paragraph of Section 2. To improve the presentation, we set the definition in its own equation.
• The grammatical error is corrected now.

Questions:

• In our experiments, Alg 1 needs less than 50 iterations to converge. In the revised version, we show the convergence of the objective over iterations to our $\varepsilon$ sensitivity study ($\varepsilon$-Experiment).
• Unfortunately, there is so far no theoretical guarantee about the convergence of Alg 1, which we hope to find in our future work.

We hope that we met all your concerns in the revised manuscript.

Many thanks for your positive evaluation of our manuscript as an 'interesting submission' that 'is theoretically comprehensive and conceptually enriches many prior OT-based distances' and for your constructive feedback! We revised our manuscript according to your suggestions. Please find our answers below.

Experiments: We agree that the shown experiments are mainly proof-of-concepts. To improve the numerical part, we extend Section 6.2 by a CPU runtime comparison with respect to the different methods. Moreover, we extend the supplement by a numerical sensitivity analysis with respect to the relaxation ($\lambda$-Experiment) and regularization parameters ($\varepsilon$-Experiment) as well as with respect to the discretization of the joint space ($Z$-Experiment). Please also note the domain adaptation experiment in Appendix C. Thanks for pointing out possible future applications which, at the moment, are beyond the scope of this paper.

Complexity: In comparison to GW, the complexity of our method is increased since we rely on a multi-marginal OT problem. However, please note that GW gives only a correspondence between the measures, but no direct information about a joint embedding. As mentioned above, we extend the feature space embedding experiment to compare the runtime between the proposed approach and other methods in the literature. Moreover, we now study the runtime in dependence of the discretization of $Z$ for the 2d embeddings in Section 6.1 ($Z$-Experiment).

Structure:

• Comment on page 1: With the lack of invariance (with respect to shapes), we mean that the computed Wasserstein plans depend on the actual embeddings of the shapes. In contrast, the Gromov-Wasserstein plans are always independent of the embeddings or representations of the considered shapes. In the revised version the comment is accordingly clarified.
• Figure 3: The figure illustrates the multi-marginal plan $\alpha$, which is, in fact, four-dimensional. The labeled marginals $\gamma_1$, $\gamma_2$, and $\pi$ give the relation to the optimization problem behind the original EW formulation. The colors match the coloring of Fig 2, which means that everything related to $\mathbb{X}_1$ is red and to $\mathbb{X}_2$ blue. We add this in the captures of the figures.
• Lastly, we will look into improving the structure and writing for the final manuscript.

Questions: The reason why our embeddings are not centered is that we show the push-forward to the complete chosen grid $Z$. For JMDS, we obtain two point clouds and manually select a region containing both, which leads to a 'centered' visualization. In fact, our and JMDS embeddings are invariant to (joint) translations.

We hope that we have addressed your concerns.

Many thanks for your positive and constructive feedback! We have revised our manuscript according to your suggestions. Please find our detailed answers to your questions below.

GW Comparison: If we consider the experiment behind Fig 1, then we are looking for rotations (isometries) which aligns the shapes/measures on the sphere/torus. A classical GW plan gives only a correspondence between the measure supports, but no direct information how to rotate the whole sphere/torus such that the supports are aligned. To say it briefly, GW cannot directly solve the considered task. Although not directly related to the embedding problem, GW and EW are compared in Fig 9 (Appendix B) for 2d shapes from FashionMNIST. Here the values of both distances strongly correlate.

Parameter Selection: To study the impact of $\lambda$, we extend the experiments in Section 6.1. The results have been added to the supplement ($\lambda$-Experiment). Figuratively, small $\lambda$ leads to fixed support barycenters, whereas large $\lambda$ gives the best possible embeddings even when no isometries exist. The two limiting cases $\lambda \to 0$ and $\lambda \to \infty$ are also supported theoretically. The proofs, which rely on similar techniques than Prop 3.4, are added to the supplement.

High-Dimensional Features The dimension of the feature space only affects the computation of the input distance matrices and not our algorithm directly. However, the size of the dataset and the chosen discretization of the joint space $Z$ directly affects the runtime ($Z$-Experiment).

Bi-Convex Relaxation: Unfortunately, we have no theoretical guarantee that the proposed algorithm converges. We choose a small, but numerically stable $\varepsilon$. In practice, we usually observe convergence after less than 50 iterations. However, in rare cases, the algorithm can get stuck in local minima. We add this discussion and a convergence experiment in the revised supplement ($\varepsilon$-Experiment).

Clustering: Please notice that Fig 6 serves mainly as an illustration of how the datasets are embedded and depends on the employed input distance. Looking at the color coding, we see that the classes are located in the same areas. Although the actual visual separation is not as clear as for JMDS, the classification accuracy is increased. To improve the visual embeddings, we now equip the feature spaces with distances from the unsupervised SCOT distance estimation routine (Demetci, et al, doi:10.1089/cmb.2021.0446). The updated distances lead to the following new version of Figure 6.

We hope that we have addressed your questions.