Joint Metric Space Embedding by Unbalanced Optimal Transport with Gromov–Wasserstein Marginal Penalization
Authors: Florian Beier
Moritz Piening
Robert Beinert
Gabriele Steidl
Venue: ICML 2025 poster (2025)
Keywords: embedded Wasserstein distance
Gromov-Wasserstein distance
optimal transport
manifold alignment
relaxed embedding
Primary area: general_machine_learning
TL;DR: We employ Gromov-Wasserstein regularization for near-isometric joint embeddings on arbitrary metric spaces.

Abstract: We propose a new approach for unsupervised alignment of heterogeneous datasets, which maps data from two different domains without any known correspondences to a common metric space. Our method is based on an unbalanced optimal transport problem with Gromov-Wasserstein marginal penalization. It can be seen as a counterpart to the recently introduced joint multidimensional scaling method. We prove that there exists a minimizer of our functional and that for penalization parameters going to infinity, the corresponding sequence of minimizers converges to a minimizer of the so-called embedded Wasserstein distance. Our model can be reformulated as a quadratic, multi-marginal, unbalanced optimal transport problem, for which a bi-convex relaxation admits a numerical solver via block-coordinate descent. We provide numerical examples for joint embeddings in Euclidean as well as non-Euclidean spaces.

Joint Metric Space Embedding by Unbalanced Optimal Transport with Cromov-Wasserstee Marginal Penalization
Florian Beier"
Moritz Piening"
Robert Beinert Gabriele Steidl
Abstract
We propose a new approach for unsupervised alignment of heterogeneous datasets, which maps data from two different domains without any known correspondences to a common metric space. Our method is based an unbalanced opti- mal transport problem with W marginal penalization. It can be seen as coun- terpart to the recently introduced joint multidi- mensional scaling method. We prove that there exists a minimizer of our functional and that for penalization parameters going to infinity, the cor- responding sequence of minimizers converges minimizer of the so-called embedded Wasser- stein distance. Our model can be reformulated as quadratic, multi-marginal, unbalanced op- timal transport problem, for which a bi-convex relaxation admits a numerical solver via block- coordinate descent. We provide numerical exam- ples for joint embeddings Euclidean well non-Euclidean spaces.
Original
Aligned
Original
Aligned
(a) Spherical embedding
(b) Toroidal embedding
Figure I. Joint (aligning) transfer of two metric spaces (gray sur- faces with surface distance) to fixed reference space, namely to sphere the torus by our method, where color "yellow" corresponds to higher values, see Subsection 6.1.
canonical joint metric space, this metric may not be suited to accurately gauge differences in their samples. Examples of such heterogeneous settings are, e.g., the comparison of graph- or mesh-valued data such as 3d shapes or mani- folds. This paper introduces a novel framework based on OT which enables the joint comparison and visualization of heterogeneous datasets by optimally transferring them into a-priori fixed metric space, see Figure for an illustrative example.
1. Introduction
The comparison of heterogeneous data distributions is a fundamental task in computer vision, computational biol- ogy and machine learning. Most existing approaches rely on using suitable ground cost function such as avail- able metric. A classic example is the Wasserstein distance which seeks an optimal transport (OT) between two given distributions. However, often the given distributions are in heterogeneous spaces, where a readily available ground cost function between these spaces does not generally ex- ist. Additional effort may be required learn appropriate cost functions (Cuturi & Avis, 2014; Heitz et al., 2021) However, even if there exists a natural embedding into
1.1. Previous Work
As first step, we highlight the most relevant contributions related to our study among the vast literature on OT and dimensionality reduction.
Optimal Transport with Invariances Classical Wasser- stein distances allow to compare measures on a common metric space. In specific tasks in geometry processing like shape matching, the considered measures, however, live on distinct metric spaces that incorporate the different geodesic distances the given shapes. To address this, Gromov- Wasserstein (GW) distances are introduced (Mémoli, 2011; Sturm, 2006). (Alaya et al., 2022), an approximation of GW distances is obtained by jointly embedding mea- sures into Euclidean spaces. Furthermore, projection- and
Equal contribution 'Institut für Mathematik, Technische Uni- versität Berlin, Germany. Correspondence to: Moritz Picning <picning @ math.tu-berlin.de.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
subspace-robust Wasserstein-2 distances are introduced in (Paty & Cuturi, 2019). Another approach incorporates in- variance to Euclidean isometries via the Wasserstein Pro- crustes problem (Grave et al.. 2019). This is extended to account for linear operators with bounded Schatten norms in (Alvarez-Melis et al., 2019) and to Gaussian mixture applications in (Salmona et al., 2024). As outlined in Sec- tion 5, our paper extends such invariant OT to non- non-Euclidean domains.
joint embeddings of the inputs into a metric space, while enabling an optimal comparison in the Wasserstein distance. 2. We prove that our functional has minimizer. Further, if its regularization parameter goes to infinity, our functional approaches the so-called "embedded Wasserstein distance" (Salmona et al., 2024). In this sense, we also refer to our model as the "relaxed embedded Wasserstein distance". 3. For the (approximate) computation of minimizer, we provide equivalent formulation of our model quadratic, multi-marginal, unbalanced OT problem. A bi- convex relaxation enables the application of existing algo- rithms to solve the problem numerically.
Joint Dimensionality Reduction Dimensionality reduc- tion a core topic in machine learning enabling visualiza- tion and clustering by finding optimal fow-dimensional rep- resentations of data. Classical methods like principal com- ponent analysis (PCA) (Greenacre et al., 2022) and multidi- mensional scaling (MDS) (Carroll & Arabie, 1998) preserve large variations or pairwise distances, but fail on nonlinear manifolds (Alaya et al., 2022; Deng et al., 2024). Nonlin- car approaches, c.g., locally linear embedding (Roweis & Saul, 2000), probabilistic models, e.g., t-distributed stochas- tic neighbor embedding (t-SNB) (Van der Maaten & Hin- ton, 2008), and deep learning methods, c.g., variational autoencoders ( (VAEs), (Kingma al., 2019) address these challenges. As an extension, several recent methods fo- cus on joint embeddings of heterogeneous data. Here, we are given data on two incompatible domains and are inter- ested in simultaneous embedding. The allignng generative adversarial network (Amodio & Krishnaswamy, 2018) employs a generative adversarial network for domain alignment. Maximum mean discrepancy (MMD) manifold- alignment (Liu et al., 2019) balances an MMD and a distor- tion term. UnionCom (Cao et al., 2020) leverages the gener- alized unsupervised manifold alignment (GUMA) (Cui et al., 2014). Single-cell alignment with optimal transport (SCOT) (Demetci et al., 2022) and the partial manifold alignment algorithm (Cao et al., 2022) employ GW distances. Fi- nally, the recently proposed joint multidimensional scaling (JMDS) (Chen et al., 2023) algorithm combines MDS with OT for the joint embedding of two datasets into a shared Euclidean space. Complementing this and research on non- Euclidean embeddings (MeInnes et al., 2018; Deng et al., 2024) and connections between MDS and GW (Van Assel al., 2024; Clark et al., 2025), we propose a joint embed- ding into arbitrary metric spaces based factored transport plans (Forrow et al., 2019) and GW distances.
4. We recall JMDS from the point of view of our model: while we are searching for the weight of atomic measures fixing their supports, JMDS fixes the weights and aims to find the supports of the measures.
5. Numerical experiments demonstrate the potential of our method for joint embeddings of heterogeneous data on the 2d Euclidean space, the torus, the 2-sphere and the space of 2d Gaussians with the Wasserstein distance.
All proofs are given Appendix A.
2. OT-Based Distances
Given a compact metric space (Z, dz) we denote by P(Z) the space of probability measures defined on the Borel-o- algebra induced by the distance dz. For two compact metric spaces 21 and Zz the push-forward measure of μ E P(Z1) under a measurable map T: Zz + Zz is denoted by
Typ = μ T E P(Z2).
The set of transport plans between two measures μ1 P(Z)) and μ22 E P(Zz) is given by
11(M, p2) = { E P(Z) x E2) : P.STT pi, 1.2).
where P;: 7, x 7/2 + 2, (21, + 24, i= 1,2, denotes the projection to the first and second component, respec- tively.
To lift dz from Z to P(Z), we rely on the Wasserstein(-2 22) distance between μl, H2 E P(Z) defined by
(ʃ LZ x²) d)(()) à (1) Zxz
W(11, = min (E)l(a
By [o(H1, we denote the set of minimizers in (1). The Wasserstein distance metricizes the weak convergence of measures, where a sequence of measures (Hu)neN con- verges weakly to a measure μ E P(Z), written pn, p. if for every (bounded) continuous function :o: 7 - IR we have Jz +p duu, + Jxx p du as + co.
1.2. Contribution
Our main contributions are the following:
1. Towards the aligned "embedding" of heterogeneous met- ric spaces into fixed, not necessarily Euclidean space, we propose to minimize an unbalanced OT problem with a quadratic cost function, where the marginals are penalized by GW distances. This formulation seeks near-isometric
To compare heterogeneous data via their internal geometry, we can rely on GW distances, which enables us to compare
2

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
measures on different metric spaces. To highlight the con- nection between measures and underlying spaces, we con- sider metric measure spaces (mm-spaces), which are triples X = (X, dxxx) {) such that (X, is a compact metric space and Ç E P( Two mm-spaces x = ,x,,,,,,,,, i = 1,2, are isomorphic if there exists a bijective map I: suppÉ1 + supp &2 between the supports of the mea- sures such that I is ceuusure-presevinng i.c. 4&1 = ε2, and 7 is an isometry, i.e., dx, (x1,
π Z, Z
X
12
π
X1 X2
X2
(a) Sturm's 's GW
(b) Mémoli's GW
dx, (a)))), (xx)) for SI By [X] we denote the equivalence class of all mm-spaces which are isomorphic to X. Finally, if / is an isometry between two metric spaces X,, i = 1, we write I: X1 + X2, and just X1 + X₂ if such an isometry between X1 and X2 exists.
Figure 2. Illustration of GW formulations of Sturm and Mémoli.
all 01,21
1, 2, be two mm-spaces, whose measures' support can be isometrically embedded into (Z, dz), i... supp + Z.
For this specific setting, we relax Sturm's GW distance to
Sturm's GW distance (Sturm, 2006) between two mm- spaces X2, = 1,2, is defined by
EW (X,, X2) =
inf 45,,
(2)
I1: supp SuppE1-2 -1 12: supp
SX, X2)
inf (Z.dz) metic space. I E
On a certain subspace of equivalence classes of mm-spaces, this defines a metric, which we call embedded Wasserstein metric. For the specific case of Fuclidean spaces X1, X2, and Z, this reduces to the Wasserstein Procrustes problem (Grave al., 2019) and to the "embedded Wasserstein met- ric" in (Salmona et al., 2024). We adopt this name.
where the Wasserstein distance is taken over the respective space (Z, dz). The SGW distance seeks optimal isometric embeddings of either mm--pace into joint metric space such that their Wasserstein distance in the embedded space is minimal, see Figure 2a. Indeed, SGW is a metric on the equivalence classes of mm-spaces.
Proposition 3.1. EW defines a metric on the subset of iso- morphic classes [(X, dx, () for which there exist surjective isomorphism I: supp + Z.
Mémoli proposed different GW distance, which numer- ically more appealing than Sturm's is construction. For two mm-spaces X,, i = 1,2, Mémoli's ' GW distance (Mémoli, 2011) is defined by
By the following proposition, the infimum in (2) is attained.
Proposition 3.2. Let X, = 1,2 be two mm-spaces and (Z,dz) a metric space such that supp + Z for 1,2. Then the infimum in (2) is attained.
GW GW X2),
min (El((().a))
Gx)
While EW relies on appropriate isometries, the following relaxation enables us to handle arbitrary mm-spaces X,, 1, 2. For À >> O, we consider the following GW penalized unbalanced OT problem
with the quadratic GW objective
Gx,xx(7)== = /f (dx; dx, xb2)
(xixX,)² x ddy(#1, (b)))
(1
EW (X,, TEP(ZxZ) 7x7 L(zzz2) d
inf
The problem seeks a transport between X, and X2 which minimizes the overall pairwise distance distortion, see Fig- ure 2b. Both distances- SGW and GW- define a metric on the equivalence classes of mm-spaces. Furthermore, both distances induce the same topology. but SGW turns these equivalence classes into a complete metric space, while GW does not (Sturm, 2006; Mémoli, 2011). The later can be alleviated by embedding into larger space and extending the GW definition accordingly (Sturm, 2023).
(Z, dz, (Z, P. P.))))) 7 (3)
By penalizing the GW terms, the marginals P..TT P(Z) take a form that enforces (Z, dz, TT to be nearly iso- morphic to the inputs x, i = 2. Furthermore, the first term ensures that the marginals P. i = 1,2 are close in the Wasserstein distance on (Z, dz). The penalization extends EW to non-isomorphic metric spaces by consid- ering dinimum-ddistoriae embeddings. By the following proposition, the infimum in (3) attained.
3. Unbalanced OT With GW Penalization
Throughout this section, let (Z, dz) be a fixed compact metric space. Furthermore, let X, (X,, Et),
Proposition 3.3. Let X, = 1, 2 be two mm-spaces and (Z, dz) a metric space. Then (3) admits a solution.
3

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
As the GW penalization enforces isometry, EW in (2) be- comes the limit of 2 in (3) if X goes infinity.
Z
Z1
Z2 = Z
Proposition 3.4. Let X i = 1,2 two mm-spaces and let (Z,dz) be a metric space such that supp & + Z. Let (An)nEN be sequence with )a sequence (TT) of minimizers of EW A. (X1, K2) verges weakly, up to subsequence, to some E PPZ Z). There exist isometries (11, 12) realizing EW (X1, X2) such that x E II
π
71
Y/2
- 3C. Then any
6
X
X2 2
Figure 3. Illustration of our multi- multi-marginal transport problem. The colors in line with Figure 2: quantities with respect to X1 red, and quantities related to X2 are blue.
If there exist no isometric embeddings supp + Z, the limit of FWW as À + oo yields best possible approxima- tions of x; on the given metric space (Z, d2) with respect to the GW distance. More precisely, a GW approximation of an mm-space on the metric space (Z, dz) is minimizer Ç E P(Z)(o of
in (3) can be rewritten
( = inf
inf
=EP(ZxZ) o ETT 1=1,2
f
GWA(X)= = min GW?(X, (2,d2,)) (EP(Z)
(4)
+ Gx,,2(12))
(6)
For discrete mm-spaces, GW approximations are studied in (Clark et al., 2025) and are closely related to MDS.
or equivalently as
(x1,
inf mL.NEE((2) rell(ar E
inf
inf
Proposition 3.5. Let X, i i= 2, be two mm-spaces and let (Z, dz) be a metric space. Let (Nnhee be a sequence with An >> oo as n. - co. Then any sequence (Nn) of minimizers of EWW,, X., (X1, X2) converges weakly, up to a subsequence, to π E II({I, (2), where S E PZ)) GW
I (2)(, M(Gx.,2(1) + Gxa,z(1a)).
approximation of x..
For computing EWA, A. we will rewrite the term by 4-plans. To this end, we use the notation Zi Z2 Z and de- note the projection onto X,.Z, by Px, Pz, respectively. Now we consider 4-plans a E PP Z, Zz X2) fulfilling Px E,, 1,2. Let
On the other side, if X + 0, the influence of the Wasserstein term in EW, x increases. Figuratively, the marginals PLa and 7T of the minimizers of EWx x have become more similar if À goes zero. In the limit case, when both marginals coincide, we obtain a fixed-support GW barycenter (Beier et al., 2023). In detail, a fixed-support GW barycenter on the metric space (Z, dz) is a minimizer Ç E P(Z)(o)) of
== Paixaa and == Px i 1.2. (8)
Clearly, such plans automatically fulfill
Pz. P2. i 1,2
see Figure Thus, (6) can be reformulated as quadratic, malti-marginal, unbalanced OT problem
2
XWB(X,,X2) = min (EP(Z)
Σ (2, (5)
FW (X,,
inf aEP(XiXZxxZ Px,,a=_u.i=1,2
Fx(a)
(9)
Proposition 3.6. Let XT i 2, be two mm-spaces and let (Z, dz) be a metric space. Let (Nn)eeN be sequence with Ao + 0 as n + Dc. Then any sequence (EN of minimizers of E.W,, E.,,((% (2,, converges weakly, up to a subsequence, to Id)s,, where Ç E P(Z) is a fixed- support GW barycenter.
with the quadratic objective
Ex(a) = /ʃ (B(()))))))) + (z(+(,
+ Σ (dx, (2 21) dz 2)))
In the rest of the paper, we skip the integration domains of the integrals for better readability, since they are clear from
X da(211 21, 22, 22) da(x1, 21, 2/2)..).
the context. Then, by definition of the GW distance, EWA
We summarize our findings in the following proposition.

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
Proposition 3.7. If solves (9,, then its projections (8) are solutions of (6) and in particular x is solution (3).
The cost function additively decouples into three partial cost functions solety depending on two coordinates of c1- Un- balanced, multi-marginal OT formulations of this kind are treated in (Beier et al., 2022), where, relying on an entropic regularization, a multi-marginal Sinkhorn scheme is pro- posed. Mathematically, this means that, for a regularization parameter & > O, we approximate the minimizer of (11) by
Conversely, any solution of (3) can be expressed this way.
In Appendix B, we illustrate relations between the Wasser- stein distance, GW, EW and A by numerical examples.
4. Bi-Convex Relaxation
/
inf Nx..00g=6 1=122
c1 dag + KL( v),
At its core, the computation of E in (9) requires the solution of a quadratic optimization problem. Similar for- mulations appear, for instance. in the computation of the GW distance (Peyré et al., 2016; Séjoumé et al., 2021), in the murti-marginal GW setting (Beier et al., 2023), and in CO-OT (Vayer al., 2020b). All of these quadratic OT problem have in common that they can be numerically solved using a block-coordinate descent their bi-convex relaxations. In the following, we adapt this approach to our multi-marginal, unbalanced transport problem (9).
where v denotes the uniform measure on X1 x 71 x 72 x X2 and KL the Kullback-Leibler divergence. More precisely, with the Radon -ikkodýn derivative dy/ du, the KL diver- gence is given by KT,(7, u) == log( dy/ de) du if & v and KL(7,, v) == oc otherwise. In total, we may approx- imate the minimizer of the quadratic, unbalanced, multi- marginal OT problem (9) by applying Algorithm 1.
Algorithm Computation of E
For this, we decouple the minimization with respect to the inner and outer transport plan y in the double integral of (9). More precisely, denoting the inner plan by a and the outer plan by 02, we consider the bi-convex relaxation
input X, = =X,,, d i i= 1,2.
(mm-spaces)
input (Z,dz)
(finite metric space) (penalization parameter)
input >> 0
input
(regularization parameter)
inf 01,0gEP(X XX2 Fx(m, 02) bax,E 1,k=1,2
(10)
1: Initialize 01 02
0=:0=::6 == E1 @ p μ &2
where j is the uniform measure on (Z, dz)
2: while not converged do
with the bilinear objective
3: Compute as in (12)
4: Update Q2 using the marti-larginal Sinkhorn scheme in (Beier et al., 2022) with 61
(2) = /ʃ (2,, 22) d3(21; =/²)
2
Compute 22 analogous to (12)
+ A Σ (dx, (1i, x/) - dz (zi, =²))²
Update j using the multi-marginal Sinkhorn scheme in (Beier et al., 2022) with cg
i=1
x don (I1, =1, 22)) daz (z'. =i. 2/2, x22)
7: end while output a1 or az
By construction, the minimizers of (10) constitute a lower bound to the original, quadratic problem (9). Moreover, every bi-convex minimizer of the form 01 = 02 yields a minimizer of the original problem.
5. Discrete EWx and JMDS
The distance EW in (2) generalizes Wasserstein Procrustes (Grave et al., 2019) to non-Euclidean spaces, while EW,, in (3) is closely related to the JMDS model (Chen et al., 2023), In this section, we explain the relation.
In order to find numerical solution, we apply block- coordinate descent, which consists of alternatively fixing on and 022 in (10) and minimizing with respect to the other argument. Fixing 01, it remains to solve the (linear) multi- marginal, unbalanced OT problem
We consider the discrete case, where
X1 = (at 1T1 and X2 = are point sets equipped with dissimilarity distances dx, and dx, as well as measures
/ (1(4), E4 x1) daz 26,26,,)))111)
inf Px pxy im 1,2
where the effective cost is given by
E =
&
2(11,22122.443 = } 2 2/) + x Σ
1=1
/ (dx, dz(z,, =1))² dan (x1, 21, 22, 22). (12)
1=1
Discrete EW, fixes discrete embedding space
Z Zg Z = {=),...,,?")
5

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
with metric dz. Accordingly, we use measures with fixed supports
This is exactly the functional proposed as JMDS in (Chen et al., 2023). Note that the Wasserstein Procrustes model is given by
14 Σ
1 82 9
Tz ,
min TEll - Ani + Iny) QESO(d)
min (TT, D2z,2.)
J,k=1
J,Km1
where = 1,2. Since the supports of the measures are fixed, we can restrict ourselves to the weight matrices μ4; E Ams &, E A... π E (22) and %/ E [ f4)). i = 1,2 in the corresponding probability simplicies A. Then E in (7) becomes the discrete minimization problem
This is exactly the first summand in JMDSs, where the minimization over the special orthogonal group SO(d) can be skipped in (13) since (2.), i = 1,2, are invariant under orthogonal transforms.
EWa(XI, X2) = min m. xE) 1112
min
min ( i=1,2
To summarize: 1.
Our discrete E fixes the support of the marginals μ4s in (7) which results in the optimization over the weights, where JMSD fixes the weights of the μui and optimizes over the supports. 2. Fixing the support has the advantage that we can work with arbitrary metric spaces (Z, dz). while the "free support" approach of JMDS is restricted the Euclidean space. 3. The minimization problems have to be tackled by completely different opti- mization algorithms, namely an unbalanced multimarginal Sinkhorn algorithm for the block-coordinate descent in Al- gorithm I for EW, and the so-called SMACOF method combined with Wasserstein Procrustes minimizations (Chen et al., 2023) for JMDS.
(2,z)),) + XGx, (m)) + (72))
where D3.2= = (d) , ) and, for i = 1,2,
Gx,,2(9))
n.mm n.m
Σ
Σ
s.k j
(dx, - 2²))
74 To
J,A=1
JMDS aims to find point sets
21 (2)....... ................................ 2 21 C RR, 22 = (22 122)...,, 22 2 Rd
6. Numerical Results
such that the dissimilarity relations in X, are approximately preserved in Z, i = 2, while ensuring that the intermedi- ate points are optimally aligned. Here, Z, i = 1,2 2 are ex- clusively in Rd with the Fuclidean metric d(<r, v) = || - y|l- Instead of measures with fixed support, we use
Next, we provide several proof-of-concepp examples'. Ad- ditional parameter ablations are described in Supplement C.
6.1. Joint Embedding of 3d Shapes
1
m
In the first example, we exploit EW to align and embed 3d shapes-the surfaces of objects in i a joint space (Z, dz). For this, we interpret a 3d shape as mm-space X == (X,dx,E). where X is the surface, dx is the surface (or geodesic) distance, and & is the uniform measure. Practi- cally, X is parametrized by the vertices of a triangular mesh, dx is approximated using Dijkstra's algorithm (Dijkstra, 1959) on the corresponding graph, and & is chosen as the discrete uniform measure. For the joint embedding of two given (discrete) shapes X,, i = 1, 2, we compute the 4-plan a in (9) using the discretization in Section 5 and Algo- rithm 1. Since the relaxation behind EW does not yield an isometry, but only a transport plan 7/- = Px, x Z we vi- sualize the computed relaxed embeddings by the marginals Pz,: Pz.
A
022: n2
6
as well as
1 j = k,
% =
with
7/th =
0
j + k,
1
and will optimize over the supports Z₁ = 1, Then (7) becomes
JMDS,(X, X2) = min
min
2el2 : (.n.) )
D2.D2 + (Z,) + Gz(Zz)), (13)
where D2 D2, = (1/ + + شسر and, for i = 1,2,
Bended Rectangles We start by embedding an S-bended rectangle and Swiss roll (with and without a hole) into R.? Intuitively, we expect that the surfaces without holes are unrolled by the relaxed embedding behind EW2, x, since there
G((Z,)= = Σ
Ri Nx (dx: (2, 27) -
1,k=1 1 1=
Σ
-
=
httpp:://github.commm wwt.s//.ccco//eeeeoeeeeeeec Re Rel
6

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
Shapes
Ours
JMDS
Shapes
Ours
JMDS
(a) S bendedd rectangle and Swiss roll (without hole).
(b) S- S-bended rectangle and Swiss roll with hole.
Figure Embedding and alignment of S-bended rectangle Swiss into R². our method, we visualize the marginals 2 a of the computed 4- 4-plan in (9) and compare them with JMDS. In the second example (b), JMDS produces unexpected hole when embedding the S-hended surface.
+
Shapes
Ours
JMDS
Shapes
Ours
JMDS
(a) Same human shape in different poses.
(b) Distinct human shapes in different poses.
Figure Embedding and alignment of human shapes from the FAUST dataset into R.? For our method, visualize the marginals P2..:a of the computed 4-plan a in (9) and compare them with JMDS. Here JMDS tends to split some of the extremitics.
actually exists isometric embeddings. For the discretization, we choose Z equispaced 50 grid [0, 1.3]2 c IR? and d2 as the corresponding Euclidean distance and apply Algorithm with 100 and ş = 10-3 In Figure 4, the re- laxed embeddings are visualized by Pz where the point sizes represent the underlying probability masses/weights. As comparison, we also show the results of JMDS with 10 and & 10-3 for the incorporated regularized OT problem. Both methods produce representations that unroll the shapes and succeed in aligning the embeddings. JMDS, however, produces an unexpected hole during the alignment of the S-bended shape and the Swiss role with hole.
are not possible. More precisely, we embed human shapes from the FAUST dataset (Bogo al., 2014) into IR2. For this, we choose (Z, dz) as equispaced 60 grid on [0, C R2 and apply Algorithm with À 100 and E = 4- 1044-10- 4 The obtained waansorr-based embeddings are shown Figure As comparison, we apply JMDS with A 10 and ε resemble the 3d human shapes, but JMDS splits some of the extremities. Since method and JMDS can both be derived from the same functional, we additionally compare the achieved objectives, the Wasserstein distances between embeddings, and the GW distances between embeddings and input spaces. The experiment is repeated for each pair- wise combination of the first ten shapes from the FAUST dataset and different X. For comparison, the shapes
10 -4-4 The 2d representations clearly
Human Shapes Since we are not limited to isometries, we next consider an experiment where isometric embeddings

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
GW2
W
Ours (A = 1)
0.0057:10.004
0.017=0.007
JMDS (A I)
0.0071=10.007
0.029=0.010
Ours (A 10)
0.0048:1.004 0.027=0.016
JMDS (A- 10)
0.0057++0.005
0.046-0.023
Table 1. Mean and standard deviation of GW distance (between in- put shapes and embeddings) and of Wasserstein distance (between both embeddings) for joint 2d embeddings of the FAUST dataset via our method and JMDS.
Ours
JMDS
Ours
JMDS
(a) SNARESeq dataset
(b) scGEM dataset
Figure 6. Joint embedding of two feature spaces into I2' using our method and JMDS (Top: first feature space, Bottom: second feature space). Both methods align the color-caded classes.
downsampled to 35 vertices, and both methods are regular- ized with & 0.001. For our method, we set Z according to uniform x0 grid on [0, 1.3]². Results are shown in Table 1. In total, we here achieve better joint embeddings than JMDS both in terms of GW and Wasserstein.
and X₂ C E2 with specific distances dx, estimated with the unsupervised SCOT routine (Demetci et al., 2022).
Spherical and Toroidal Embeddings Finally, we con- sider the alignment of spherical and toroidal subsets, i.e., the joint embedding into a -on-Euclidean space. More precisely, we consider the 3d shapes in the introductory Figure 1. As target space (Z, dz), we choose 30 grid on the canonical parametrization of the sphere S(2) and the torus T(2) = S(1) x 8(1) equipped with the corresponding geodesic distance. Using x= = 10x and ş 10 3 for Algo- rithm 1, we embed the spherical rectangle and a spherical cap. both with a hole, onto the sphere, and the half- torf- and toroidal triangle onto the torus. The relaxed embed- dings are shown in Figure 1, where the color encodes the mass of the computed marginals. Up to smoothing due to the incorporated entropic regularization, the transport plans % Px, 20 correspond to the expected isometries.
2. We equip the constructed metric spaces with the uniform measure to obtain the mm-spaces X, = (X,, dx.,
3. Fixing a 20x20 grid L = {z .....,"") C I² over
[0. 1]2 equipped with the Euclidean metric, we apply Algo- rithm with the discretization in Section to compute the 4-plan a in (9).
4. Based on the transport plans % 2. a pair xt with its barycentric projection . E IR. given by
ZR
For the first aim, the identification of the underlying cor- respondence, we identify of with x whose barycentric projection iny is closest to मं. To quantify the quality of the pairing, we rely on the FOSCTTM (fraction of samples closer than the true match) score (Liu et al., 2019). More precisely, FOSCTTM defined by
6.2. Alignment of Feature Spaces
In the next example, we use our method to align different feature spaces occurring in real-world data. More precisely, we consider the genetyping- determination of the corre- sponding cell class of single cells from various measured modalities. Multimodal technologies as proposed in (Cheow al., 2016; Chen et al., 2019) are, however, uncommon such that there is rising interest in embedding the corresponding features into joint feature space (Demetci et al., 2022; Liu
Σ # k (1,...,,,)) n
m) |||
< |
1
and ranges from 0 (perfect identification) to 1.
For the second aim, the transfer of a classifier from X₃ to X2, we exemplarily consider the k-nearest neighbor (KNN) method. More precisely, we use the barycentric projections (a),... (.,,..., ..T) and the corresponding gene type labels of the first modality to classify feature 4. (more exact á from the second modality, where we consider the closest five neighbors. Here, higher classification accuracies (KNN- Acc) indicate better class alignments.
et al., 2019; Chen et al., 2023). Following the experiments of (Demetci et al., 2022) for n single cells, our first aim is to as- sign the recorded features X, := (=t....,..)) " C Rdi from one modality to the features X2 == {a).... .........) x²}} C R2d of another modality, i.e., recover the underlying one-to- one correspondence between æ/ and 2z.. Our second goal consists in transferring a classifier from X, to X₂ in an unsupervised manner.
For the experiments, we employ the publicly available 2atasets² from (Demetci al., 2022), where SNAREseq
To achieve both goals, we apply the following methodology: We equip the recorded, uncorrelated features X, C Rdd.
ihtrps://rsinhhlabbgg tith. http:/gggtti...oo////aaaaaa io/ScoT/data//
8

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
SNARESSQ
ICACCTTML
EW, 2013)
0.165
28.9s
0.610
(Denetcl 2003)
0402
0.578
17.66
UnionClan (Cas atal. 2000)
0496
04663
1.5
Table Comparison of the joint embedding of two feature spaces into R² using our method, JMDS, SCOT, and UnionCom.
consists of 19 and d2 10 features of 1047 sin- gle cells, and scGEM of d1 34 and d2 27 features of 177 specimens. The computed embedding of X; into I of our method and JMDS are visualized in Figure 6. For FOSCTTM and KKN-Acc, we additionally compare both methods with the alignment techniques UnionCom (Cao et 2020) and SCOT (Demetci et al., 2022) using the same distances on X, Note that UnionCom relies on GUMA and t-SNE, whereas SCOT relies on GW and MDS. The results recorded in Table 2. The hyperparameters of all meth- ods are chosen according to grid search minimizing the FOSCTTM on 10% validation split, see Supplement D. A further experiment with different feature spaces for MNIST and FashionMNISTT is given in Supplement
Original
Aligned
Original
Aligned
(a) Tri-modal GMM
(b) "Moons" GMM
Figure 7. Alignment of GMMs with respect Wasserstein distance. The considered and computed GMMs are visualized via their density function on R.²
While our model can handle the aligned transfer into arbi- trary metric spaces, it relies on their appropriate definition. In particular, due to computational restrictions, working on grid in IR is restricted to small dimensions d. As future research direction, we are interested in developing non- Euclidean free-support solver based on existing tree-support GW barycenter algorithms (Peyré et al., 2016; Vayer et al., 2020a; Beier & Beinert, 2025). This may combine our EWA approach with a non-Euclidean version of JMDS.
6.3. Alignment of Gaussian Mixture Models
In the final example, inspired by (Salmona et al., 2024), we want to align Gaussian mixture models (GMMs). For this, we consider the space of 2d Gaussians 22 = {N Σ) E IR2 E R2x2 }pd.} equipped with the Wasserstein dis- tance, which yields non-Euclidean geometry. A GMM be interpreted mm-space X = (N2, E) with discrete measure == Σラ= Ej Sy For the experiment, we fit two GMMs to affine-tran trine-transformed 2d datasets using the expectation- ciniiionnmaiiizzitiinnn algorithm. Here, we employ the datasets "Blobs" and "Moons" from seikit-learn (Pedregosa al., 2011), see Figure For the alignment of the GMMs, we choose Z 22, where the means form a 15x 15 grid over [0, 1 and we coarse grid over suitable matrices. Particularly, we consider matrices
From a theoretical point of view, the joint embedding using proposed variational formulation may adapted to handle multiple input spaces. The main issue is hereby the selection of an appropriate generalization of the employed Wasser- stein part. A possible choice would be to utilize only the Wasserstein distance between the embeddings of the kth and (k + 1)st embedding. 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.
o Σ
--
012
of
where we choose of of E (0.8.1) and 012 E (-0.2, 0., 0.2). The parameter , 0 corresponds to the mean variance of the considered GMMs. Applying Algo- rithm 1 to compute a in (9), and considering the marginals := a, which are again discrete measures on Ng. we obtain the aligned GMMs in Figure 7.
Acknowledgements
M.P. acknowledges funding from the German Research Foundation (DFG, GRK2260 BIOQIC project 289347353).
Impact Statement
7. Conclusion
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
We propose an unbalanced OT framework with GW marginal penalization which enables the joint embedding of two datasets based on pairwise intra-dataset distances.

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
References
Chen, S., Lake, B. B., and Zhang, K. High-throughput se- quencing of the transcriptome and chromatin accessibility in the same cell. Nat. Biotechnol., 37:1452-1457, 2019.
Alaya, M. Z., Bérar, M., Gasso, G., and Rakotomamonjy, A. Theoretical guarantees for bridging metric measure embedding and optimal transport. Neurocomputing. 468: 416-430, 2022.
Cheow, L. F., Courtois, E. T., Tan, Y., Viswanathan, R., Xing, Q.. Tan, R. Z., Tan, D. S. w., Robson, P., Loh, Y.-H., Quake, S. R., and Burkholder, W. F. Single-cell multimodal profiling reveals cellular epigenetic hetero- geneity. Nat. Methods, 13:833-836 2016.
Alvarez-Melis, D., Jegelka, S., and Jaakkola, S. Towards optimal transport with global invariances. In Interna- tional Conference Artificial Intelligence and Statistics (AISTATS), volume 89 of Proceedings of Machine Learn- ing Research, pp. 1870-1879. PMLR, 2019.
Clark, R. A., Needham, T., and Weighill, T. General- ized dimension reduction using semi-relaxed Gromov- Wasserstein distance. In AAAI Conference Artificial Intelligence (AAAI), volume 39, pp. 16082-16090. AAAI Press, 2025.
Ambrosio, L., Gigli, N., and Savaré, G. Gradient Flows in Metric Spaces and the Space of Probability Measures. Birkhäuser, Basel, 2005.
Cui, Z., Chang, H., Shan, S., and Chen, X. Generalized unsupervised manifold alignment. In Advances in Neural Information Processing Systems (NeurlPS), volume 27. Curran Associates, Inc., 2014.
Amodio, M. and Krishnaswamy, S. MAGAN: Aligning biological manifolds. In International Conference on Machine Learning (ICML), pp. 215-223, PMLR, 2018.
Beier, F. and Beinert, R. Tangential fixpoint iterations for Gromov-Wasserstinn barycenters. SIAM Imaging Sci., 18(2):1058-1100) 2025.
Cuturi, M. and Avis, D. Ground metric learning. JMLR, 15 (1):533-5664 2014.
Demetci, P., Santorella, R., Sandstede, B., Noble, W. S., and Singh, R. SCOT: single-cell multi-omics alignment with optimal transport. J. Comput. Biol., 29(1):3-18. 2022.
Beier, F., von Lindheim, J., Neumayer, S., and Steidl, G. Unbalanced multi-marginal optimal transport. J. Math. Imaging Vision, 11:99--413, 2022.
Deng. C.. Gao, J., Lu, K., Luo, F., Sun, H., and Xin, C. Neuc-MDS: Non-Euclidean multidimensional scaling through bilinear forms. Advances on Neural Informa- tion Processing Systems (NeurIPS), volume 38. Curran Associates, Inc., 2024.
Beier, F., Beinert, R., and Steidl, G. Multi-marginal Gromov W Wasserstein transport and barycentres. Inf. In- ference, 12(4):2733-7781 10 2023. ISSN 2049-8772.
Bogachev, I. Weak convergence of measures. Number 234 in Mathematical Surveys and Monographs. American Mathematical Society, Providence, RI, 2018.
Dijkstra, F. W. note on two problems in connexion with graphs. Numer Math., 1:269-271, 1959.
Bogo, F., Romero, J., Loper, M., and Black, M. J. FAUST: Dataset and evaluation for 3D mesh registration. In IEEE Conference Computer Vision and Puttern Recognition (CVPR). IEEE, June 2014.
Forrow, A., Hütter, J.-C., Nitzan, M., Rigollet, P., Schiebinger, G., and Weed, J. Statistical optimal trans- port via factored couplings. In International Conference on Artificial Intelligence and Statistics (AISTATS), Pro- ceedings of Machine Learning Research, pp. 2464-2465. PMLR, 2019.
Cao, K., Bai, X., Hong, Y., and Wan, L. Unsupervised topo- logical alignment for single-cell multi-omics integration. Bioinformatics, 36(1):i48-i56, 2020.
Grave, E., Joulin, A., and Berthet, Q. Unsupervised align- ment of embeddings with Wasserstein procrustes. In Chaudhuri, and Sugiyama, M. (eds.), International Conference on Artificial Intelligence and Statistics (AIS- TATS), volume 89 of Proceedings of Machine Learning Research, pp. 1880-1890, PMLR, 2019.
Cao, K. Hong, Y., and Wan, I.. Manifold alignment for heterogeneous single-cell multi-omics data integration using Pamona. Bioinformatics,
2022.
Carroll, J. D. and Arabie, P. Multidimensional scaling. In Birbauu, M. H. (ed.), Measurement, Judgment and Decision Making, pp. 179-250 Elsevier, 1998.
Greenacre, M., Groenen, P. J., Hastie, T., d'Enza, A. I., Markos, A., and Tuzhilina, E. Principal component anal- ysis. Nat. Rev. Methods Primers, 2(1):100, 2022.
Chen, D., Fan, B., Oliver, C., and Borgwardt, K. Unsu- pervised manifold alignment with joint multidimensional scaling. In International Conference on Learning Repre- sentations (ICLR). OeenReview, 2023.
Heitz, M., Bonneel, N., Coeurjolly, D., Cuturi, M., and Peyré, G. Ground metric learning on graphs. J. Math. Imaging Vision, 63:89-107, 2021.
10

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
Kingma, D. P., Welling, M., al. An introduction varia- tional autoencoders. Found. Trends Theor. Comput. Sci., 12(4))307-392, 2019.
Séjourné, T., Vialard, F.-X., and Peyré, G. The unbalanced Gromov-Wasserster distance: Conic formulation and relaxation. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems (NeurIPS), volume 34 Advances in Neural Information Processing Systems. Curran Associates, Inc., 2021.
Liu, J., Huang, Y., Singh, R., Vert, J.-P., and Noble, W. S. Joindy embedding multiple single-cell omics measure- ments. 19th International Workshop on Algorithms in Bioinformatics (WABI 2019), volume 143 of Leibniz International Proceedings in Informatics (LIPlcs). pp- 10:1-10:13, 2019.
Van Assel, H., Vincent- C., Courty, N., Flamary, R., Frossard, P., and Vayer, T. Distributional reduction: Unifying dimensionality reduction and clustering with Gromov-Wasserser arXiv preprint arXiv:2402.0223)) 2024.
Mcinnes, L., Healy, J., Saul, N., and Großbberger, L. UMAP: Uniform manifold approximation and projection. J. Open Source Softw., 3(29), 2018.
Van der Maaten, L. and Hinton, G. Visualizing data using t-SNE. JMLR, 9(86):5579-200 2605, 2008.
Mémoli, Cromov-Wassersee distances and the metric approach to object matching. Found. Comput. Math., (4):417-487 2011.
Vayer, T., Chapel, L., Flamary, R., Tavenard, R., and Courty, N. Fused Stomov-Wasserstee distance for structured objects. Algorithms, 13(9):212, 2020a.
Mordukhovich, B. S. and Nam, N. M. Convex analysis and beyond. Springer, 2022.
Vayer, T., Redko, Flamary, R., and Courty, N. Co-optimal transport. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. and Lin, H. (eds.), Advances in Neural In- formation Processing Systems (NeurIPS), volume 33 of Advances in Neural Information Processing Systems, pp. 77559-17570. Curran Associates, Inc., 2020b.
Munkres, J. Topology. Featured Titles for Topology. Pren- tice Hall, Upper Saddle River, 2000.
Paty, F.-P. and Cuturi, M. Subspace robust Wasserstein dis- tances. International Conference on Machine Learning (ICML), pp. 5072-5081. PMLR, 2019.
Villani, C. Topics in optimal transportation, Number 58 in Graduate Studies in Mathematics. American Mathemati- cal Society, Procidence, 2003.
Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., al. Scikit-learn: Machine learning in Python. JMLR, 222282522830 2011.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-MNIST: A novel image dataset for benchmarking machine learning algorithms. arXiv preprint (eXXi:1700.077700 2017.
Peyré, G., Cuturi, M., and Solomon, J. Gromov- Wasserstein averaging of kernel and distance matrices. In Balcan, M. and Weinberger, K. Q. (eds.), International Conference on Machine Learning (ICML), volume 48 of Proceedings of Machine Learning Research, pp. 2664- 2672. PMLR. 2016.
Roweis, S. T. and Saul, I.. K. Nonlinear dimensionality reduction by locally linear embedding. Science, 290 (5500))2323-2326, 2000.
Salmona, A., Desolneux, A., and Delon, Gromov- Wasserstein-like distances the Gaussian mixture mod- els space. TMLR, 2024.
Sturm, K.-T. On the geometry of metric measure spaces. Acta Mathematica, 196(1):65 131, 2006.
Sturm, K.-T. The space of spaces: curvature bounds and gradient flows on the space of metric measure spaces, volume 290 of Memoirs of the American Mathematical Society. American Mathematical Society, Providence, 2023.

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
A. Proofs
The main idea of the proof follows the argumentation for the original embedded Wasserstein distance on Euclidean spaces in (Salmona al., 2024), which we adapt to arbitrary metric spaces. The definiteness follows directly from the definition of the isomorphic equivalence classes, the isometries in EW, and the definiteness of the Wasserstein distance. The symmetry is inherited from the Wasserstein distance too. It remains show the triangle inequality. Without loss of generality, we take representatives X; == (2, dz, Et), 10 1,...,3. 3, such that & E P(2) possess full support. Since the inverse and composition of isometries are isometry, we have
Proof Proposition 3.1
(W, X2)
inf WI,,,Si,, 12,522) TuZ=Z
inf W(11, (T 12)x52) inf W(f,, 4,22 I=Z=Z
1Z-+
Using the triangle inequality of the Wasserstein distance, we obtain
EX(XI,
inf W(S, 1422) V inf J451) + W(J4aa, T252) 1,3:4+2
1:2-2
(W(f1, + inf (U4S):1$62)) J:2-2
inf
inf W(11, + inf W(fs, (J-1 o 1)ta)) 1:242
1:2-2
J:2-2
inf W(11, J3Ea) + inf J.ZIZ
W(f,, K2ł2) I (3)) EW (X3,
kiZ+Z
Proof of Proposition 3.2
Let X E (X; : = 1.2]. We equip the space of continuous functions C(X, Z) with the topology of uniform convergence d(f.g) = max,ex dz (f((()))(x))
First, we we show that
Iso X, =) (T: 2J C 2(
compact. Since (Z, dz) is compact, we see immediately that Iso X, Z) is pointwise bounded. Moreover, since all functions in Iso(X, Z) are isometries the set is equicontinuous. By the Theorem of Arzelà-Ascoli (Munkres, 2000) (Thm. 47.1), this gives the relative compactness of Iso. Further, Iso 2) C(X, Z) is closed and thus compact, since have for any sequence of isometies (In), from X Z converging to some I C(X, Z) that
(((()))))) + d2 I(x))
as n + oo, so that / E Iso Z). Then also (so(X1, 2) x 1X,, 2) is compact. Next, we show that the functional
F(I, 12) = ((,,,E11, 12.5
in (2) is lower semi-continuous on Iso (6(X1,2) x Iso [x,, 2). For i = 1, let (JIn))eN C Iso(X, Z) and I; E Iso Z) with Ii,n >> T as n > co. Then, for all y E C(Z, RR, it holds by the dominated convergence theorem that
/ 4 d((km)) = / go po
/ X.
4o I, -f pdisss as n 00.
=/pdlisss
Thus we obtain the weak convergence (Jem)44 Ta. Finally, due to the joint weak lower semi-continuity of the Wasserstein distance (Ambrosio et al., 2005) (Lemma. 7.1.4), we get the desired lower semi-continuity of F.
Finally, application of the Weierstrass theorem yields the assertion.
Proof of Proposition 3.3
Since P(Z x Z) is weakly compact, the existence of minimizer follows by Weierstrass' theorem, see, e.g., (Mordukhovich & Nam, 2022) (Thm. 2.164), once show that the objective function
F(T) = f B(2,z Jzxz
Σ
GW2(X,) (Z, P25
in (3) weakly lower semi-continuous,
12

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
Let (Tn)NEN C P(Z x Z) and P(Z 2) be such that Tn π as n cc. We handle the terms of F separately. Since d3 is continuous in both arguments, obtain by definition of the weak convergence that
[ Jaxz 2') dTTl(z2)) - / (2 dz(
2x7
0o. We turn to the GW terms. For i 1,2, set
Pi, == Tn and 14 == PI.
Then the weak lower semi-continuity of the marginal projection operators yields Pi, H4 as + oc. Now let 7i,m E [S (i,n) be solution of GW,,, (Z, dz, (i,n),, n E N. As (Min)DeN P(X; 2) and the latter is weakly compact, choose converging subsequence again denoted by (yi, N)ENN such that Vi % E P(X, Z). Using again the weak continuity of the marginal projection, we get E IT( Therefore, it holds
GW (Z, P, GW2 GX, (2,d (Z,,,,,,) < // (dx,( (1
(Xxx2)²
and by continuity of the integrand and the fact that (7im 8 Ti,)) - (7i 60 7:), see, e.g. (Bogachev, 2018) (Prop. 2.7.8) further
GX, P. s lim // (dx, 24) dz (2(,2'))) dy (1,, 2) (x),(()),)))
n-
(X,x2)2
lim GW2 (Z, d2, Mi,,)) n-+00
lim GX, (Z, dz. p.)
n
This gives lower semi-continuity of the functional of F.
Proof of Proposition 3.4 1. First, we show that
EW E X2) & EW(X,,X))
A> 0.
(14)
To this end, let 4: X; + Z; be isometries realizing (2)), i.e.,
EW(X,,X,) 12,15
By definition of the GW distance, we have GW (i,,)) = 0, = 1,2. Then we obtain for ē E 1o 12t52)
that
EW XWA,X, S (/ d2 di(z,z1))
EW (Z,d2) (X,, X2). = W(1,, (((2,2)
2, 2,x22
Now let (Am)DEN with An 00 as n + oc, and let nr E P(Z Z) realize EW x,, (X,, X2). Then we conclude by (14)
that
-
CW, X2) 2 EWA, (2),
dz(z, 2') + An Σ (2,7,
1xZ
2
2 An Σ
GW(X,, (Z, (T:)
1=1
Hence, we obtain for i = 1, that
GW (2, d2, (P)) + 0 as n +
Since (NT) is contained in the weakly compact P(Z Z), there exists a subsequence (denoted in the same way) which weakly converges some π E P(Z x Z). Then also TT converges weakly P. i 1, 2 and we obtain as in the proof of (3.3), that
GW (Z, dz, Pa.T)) < lim GW (2, dz, P.ςTT))) = 0.
400
13

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
Thus, there exist isometries 1, X3 Z with P. 1,55 i,,, a E (1 12,52). Finally, obtain for any
fixed À that
1,,, (1/2)//000 22)
(fanea 22 +
GW(X,, (Z, )1
>
< lim
dT, 22) +
GW(X;, (Z, dz, Pig
1=00
< lim EWA, (2), (W X2). m000
Hence ( 12,552) and
W(Jrski, = (1 ((()) dT(((,)
Xx)
inf 22,112
implies that (4, I2) minimizes EW ( X2) and x realizes W (2.15251
Proof of Proposition 3.5 Without loss of generality, assume that nn converges to II( (2) with S, P(Z) weakly; otherwise take convergent subsequence. Since (Z, dz) is compact, the diameter M == s' d(z, z') is finite. Using plans 7n E II(SI,n, (2,n), where Si.. E P(Z) are arbitrary GW approximations of X,, i.c., minimizers of (4), we estimate the relaxed embedded Wasserstein metric as follows:
(2, dz, PI. GW dz, P2. S An
F.W (X1, X2)
/
M
$ An Jzxz
+ GX,, dz, Pag S
+ GWA(XI) + GW)(X.).
An
Taking the limit + co, and exploiting (Beier & Beinert, 2025, Lem. I.1,, we notice that the left-hand side becomes GW (2, d2, (1) + GW, (Z, dz, (2)). Hence C,, i = 1, 2, have to be GW approximations of X, on (Z, d2). O
Proof of Proposition 3.6 Without loss of generality, assume that n. converges π E P(Z 2) weakly; otherwise
take convergent subsequence. Estimating EW, x., using the plans म. (Id, Id) where Sn P(Z) is an arbitrary fixed-support barycenter, i..., minimizer of (5), we obtain
FW2 (2) S Am GW, (Z, Pa.ãn)) A, GWB(().
Thus, An implies EW,, x., (X,,X, + and therefore W 1n, 0. Due to the stability of the Wasserstein distance (Ambrosio et al., 2005, Prop. 7.1.3.), the limit of n. has the form [d,, for some Ç P(Z). Moreover, we have
EW X2) S GWB(X,,x,)) GW (x,, P(,5,) GW (X,, p2) S X,
Taking the limit co, and exploiting (Beier & Beinert, 2025, Lem. I.1), we notice that the left-hand side becomes GW (Z, dz ()) GX, (Z, dz, ()). Hence Ç has to be fixed-support GW barycenter.
14

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
Proof of Proposition 3.7 For arbitrary a E P(X1 x Z Zo x X2), and for π P2i x 2 and 7 Paixz.da, the functional in (9) can be estimated by
Fx(o) - // &2 + +
(x, dz(z,; da( 21, #1, 22) da(w z 222 2/22)
(x1x2122xXx1)
- 12 2)) 22) + Σ/ //
2
(X1x2))2 (dx,(x;;))) dz(z,, 23)) dril 2i) dn((d/
= x,22 (7 22) + A (G,,, + Gxaz((a)) (x1,
(15)
Now let x² and 7f' be solutions of (6). Due to the gluing lemma (Villani, 2003, Lem. 7.6), there exists a* E P(X, x Z1 x Z2 X2) such that Pzzx 72, = x" and 2x, 7 For this a", the inequality in (15) becomes sharp, i.e., Fx EEW (X1, X2). Thus a" solves (9). The other way round, let a be a solution of (9). As argued in the first part of the proof, the inequality in (15) has to be again sharp. Hence π² = Pzi x 22, and 7,- == Px. solve (6).
B. Numerical Studies on GW, EW, and EW
Our approach enables geometrically meaningful data comparison as we can approximate the EW metric on general metric spaces Z. We validate this in two experiments.
Approximation of EW by EWa for Synthetic Circular Data Consider circular mm-spaces X, - ( E:) with (Z, dz) = (8', ), where (S', dg.) ) denotes the circle equipped with a circular distance. We set &1 uniformly distributed and ε2 according to the density of a von Mises distribution with increasing dispersion parameter K. Note that the von Mises distribution takes the form of uniform distribution for rotational invariance of the uniform distribution on the circle, the embedded Wasserstein distance can easily be calculated in this case as it coincides with the Wasserstein distance. Thus, we can ignore the isometry in EW and directly calculate EW by solving the linear program underlying the OT problem. We discretize the circle into 360 bins and estimate EWx for different choices of x and The results in Figure 8 show an improved approximation of EW for larger À as announced in Proposition 3.4. Indeed, we have an excellent fit for X 20, whereas X 0.2 leads to an underestimation.
and becomes more concentrated for k. - Due to the
2d Shape Matching Next, we compare randomly rotated gray-value images. We can describe such images as mm-spaces X, = (X,d&C, where C TR22 is grid that describes the pixel positions equipped with the Euclidean metric and & the pixel intensity, see (Beier al., 2023). We use (Z, dz) (X,d). We apply random affine transformations, i.e., translations and rotations, to the first 10 FashionMNIST (Xiao et al., 2017) training images from the classes "Trouser", "Pullover" and "Sneaker", respectively. Then we compute W, GW and EW, x for X 20. The results are displayed in Figure 9. As expected, the Wasserstein distance cannot recover the class structure due to the affine transformations, whereas FWW x and GW capture the class structure. Moreover, we see that the distance matrices of EW, and GW display almost identical patterns, up to scaling.
EW
EWT
EWE
2
8
10
6
8
10
2
4
6
8
10
K
2 (A = 00)
EW3 (A = 20)
EW? (A = 0.2)
Figure 8. Comparison of EW with EW, w for different While the curves for EW (left) and EW 20 (middle) are almost identical, those for EWo.2 (right) is consistently smaller.
15

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
T
P
P
S
S
5
S
T
P
P
EWA (A = 20)
GW
W
Figure Pairwise W. GW and E0 distances randomly rotated and translated images from each of "Trouser" (T), "Pullover" (P), and "Sneaker" (S) MNSST classes. Smaller distances are darker. The class block structure is recovered with EW,, and GW, not with Wasserstein.
X2
À 0.5
A
À 100
À 0.5
==1
100
(a) S- S-bended rectangle and Swiss with hole.
(b) Same human shape different poses.
Figure 10. Ablation study of X for joint embedding of shapes from Figure 4b and Sa. Based on varying À values, employ our method with 0.001 and Z defined by uniform 20 grid [0, 1.3/3.
C. Parameter Sensitivity
Algorithm depends on three input parameters, namely the GW regularization parameter the entropic regularization aim set À high as possible and low
parameter and the chosen discrete reference space Z. Generally, possible while preserving numerical stability. Based two sets of 3d shapes described in Section 6.1, namely s-curve in combination with Swiss roll with hole and two human shapes, we present a parameter sensitivity study that illustrates the impact of X, and Z. Throughout this study note that our Euclidean joint embeddings are invariant to translations and rotations which results arbitrarily oriented joint embeddings. Figure 10, we investigate the influence of the GW parameter that promotes near-isometric embeddings. We observe that the embeddings of both spaces nearly coincide for small X. This supported by Proposition 3.6 showing that the embeddings form a fixed-support barycenter in the limiting case. Moreover, we see saturation effect where an increase of À from to 100 has only small impact on the joint embeddings. In Figure 11, investigate the influence of the entropic parameter that enables efficient computation. Here, we see that large leads to highly blurred embedding. This shows that it is advisable to choose low, but numerically stable value for s. Lastly, we investigate the discretization of the reference space in Figure 12. Again, we aim for a fine discretization, but the runtime of our algorithm depends heavily on the size |Z| of our reference space. We observe that the joint embeddings display the same shapes and alignments, only at higher resolution. All embeddings were computed with 40 iterations Algorithm
16

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
Functonnl
Funcs
Functional
Functionel
Tuncionall
20
10
20
10
20
20
8
0.001
0.002
8 0.01
0.001
8 0.002
E 0.01
(a) S-bended rectangle and Swiss roll with hole.
(b) Same human shape different poses.
Figure 11. Ablation study of for joint embedding of shapes from Figure 4b and 5a. Based varying values, employ method with À = 100 and Z defined by a uniform 20 20 grid in [0, 1.3] Below the joint embeddings, we show convergence plots of the unregularized functional (a) from (9) after 12, 16, 20, 24 and 28 iterations.
|2| 15a (28s)
|Z| 20
[Z| 60
|Z| 152
|Z| = 20
60 (20min)
(53s)
(27min)
(19s)
(34s)
(a) S-bended rectangle and Swiss roll with hole.
(b) Same human shape different poses.
Figure 12. Ablation study of the grid size for joint embedding of shapes from Figure 4b and 5a. Based 0.001 and À 100, employ our method with Z C [0, 1.3)² defined by a uniform 15 x 15, 20 x 20 or 60 60 grid. We include CPU runtimes.
D. Hyperparameter Selection
JMDS, SCOT, and our method all rely on the Sinkhorn algorithm. Therefore, we use the same fixed entropic regularization for all three methods for a given experiment. We choose it as small as possible while avoiding numerical overflow, i.c., E = 0.001 or 0.002. Next, there are multiple ways defining the nearest neighbor graph described in Section 6.2. For our experiments, we start by rumning the SCOT algorithm which offers an unsupervised hyperparameter search to determine the shape of the nearest neighbor graph (Demetci et al., 2022). We use the unsupervised hyperparameter selection of the
17

---

Joint Metric Space Embedding by Unbalanced OT with GW Marginal Penalization
publicly available SCOT implementation? to estimate suitable nearest-neighbor graph. Here, we the the rest of the experiment and use the default grid over the neighborhood size. Note that the public implementation uses connectivity graph based pairwise correlation distances between points instead of Euclidean distances, see (Demetci et al., 2022). Then, we use SCOT and MDS to estimate the final embeddings based on the pairwise geodesic distance matrix estimated from this graph. We use the same distance graph for UnionCom, JMDS, and our algorithm. Here, we estimate suitable hyperparameters based on a random 10% validation split. We use four parameter configurations for each algorithm, i.e., we test 0.01, 0.1, 10, 100 for JMDS and our method. For UnionCOM, we use combinations of / = 1, 10 and a perplexity of 30 and 100 for the t-SNE dimensionality reduction. We use no parameter annealing for any method.
in
E. Further Example for Comparing Feature Spaces
As an addition to the genetyping experiment, we present another experiment with real-world data based on latent spaces. We consider the task of comparing the 4d latent spaces of an auto-encoder (AE) versus a variational auto-encoder (VAE), trained on FashionMNIST (Xiao et al., 2017). For this purpose, we train a simple convolutional AE with two convolutional layers and one linear layer for the encoder and the decoder. We train the AE with the MSE loss and the VAE with the log-likelihoon loss for 10 epochs the canonical training data splits. Subsequently, we embed the test data using the resulting AE and VAE into the 4d latent spaces. Note that such distinct AE trained the same dataset generally produce incomparable embeddings. For our experiment, we consider the first 100 data points of the test split. We choose Z as equispaced 20x20 grid on [0, 1.3] C R² with the Euclidean distance. We follow the same evaluation routine as in Subsection 6.2 for parameter selection and visualization. The joint embeddings are visualized in Figure 13. Comparing FOSCTTM and KNN, we see again a good embedding quality for our approach, see the figure caption.
Figure 13. Joint embedding of the FashionMNIST latent space of an AE and a VAE into I² using our method and JMDS. AE latent space is on the top and the latent space is on the bottom. Both methods align the color-coded classes. A quantitative comparison shows that our model achieves better FOSCTTM (0.040) than JMDS (0.042), whereas our KNN-Accuracy (0.589) slightly worse than the one of JMDS (0.644).
Ours
JMDS
http.///tttuuucccc hattppe//gggtabt..c///////agaa// (March 2025)
18


Reviews:
[{"review_id": "Ke0BET5hKX", "reviewer": "Reviewer_cuNW", "summary": "The paper proposes an unbalanced optimal transport problem with Gromov-Wasserstein marginal penalization, which can map data from two different domains, where no known correspondences exist, into a common metric space. The problem can be resolved using block-coordinate descent based on the Sinkhorn algorithm. A series of experiments, including joint embedding of 3d shapes, alignment of feature spaces, and alignment of Gaussian mixture models, have verified the effectiveness of the proposed EW_lambda distance.\n\n## Update After Rebuttal\nThank you for the authors' response, which has resolved my confusion. I keep my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for your positive and constructive feedback!\nWe have revised our manuscript according to your suggestions.\nPlease find our detailed answers to your questions below.\n\n**GW Comparison:**\nIf we consider the experiment behind Fig 1,\nthen we are looking for rotations (isometries)\nwhich aligns the shapes/measures on the sphere/torus.\nA classical GW plan gives only a correspondence\nbetween the measure supports,\nbut no direct information \nhow to rotate the whole sphere/torus\nsuch that the supports are aligned.\nTo say it briefly, GW cannot directly solve the considered task.\nAlthough not directly related to the embedding problem,\nGW and EW are compared in Fig 9 (Appendix B)\nfor 2d shapes from FashionMNIST.\nHere the values of both distances strongly correlate.\n\n**Parameter Selection:**\nTo study the impact of $\\lambda$,\nwe extend the experiments \nin Section 6.1.\nThe results have been added to the supplement\n([$\\lambda$-Experiment](https://imgur.com/a/g1Thyrg)).\nFiguratively,\nsmall $\\lambda$ leads to fixed support barycenters,\nwhereas large $\\lambda$ gives the best possible embeddings\neven when no isometries exist. \nThe two limiting cases $\\lambda \\to 0$ and $\\lambda \\to \\infty$\nare also supported theoretically. \nThe proofs,\nwhich rely on similar techniques than Prop 3.4,\nare added to the supplement.\n\n**High-Dimensional Features**\nThe dimension of the feature space\nonly affects the computation of the input distance matrices\nand not our algorithm directly.\nHowever, the size of the dataset\nand the chosen discretization of the joint space $Z$\ndirectly affects the runtime ([$Z$-Experiment](https://imgur.com/a/rlnnkVr)).\n\n**Bi-Convex Relaxation:**\nUnfortunately,\nwe have no theoretical guarantee\nthat the proposed algorithm converges.\nWe choose a small, but numerically stable $\\varepsilon$.\nIn practice,\nwe usually observe convergence after less than 50 iterations.\nHowever, \nin rare cases,\nthe algorithm can get stuck \nin local minima.\nWe add this discussion and a convergence experiment\nin the revised supplement\n([$\\varepsilon$-Experiment](https://imgur.com/a/odRSmkl)).\n\n**Clustering:**\nPlease notice that\nFig 6 serves mainly as an illustration\nof how the datasets are embedded\nand depends on the employed input distance.\nLooking at the color coding,\nwe see that \n the classes are located in the same areas.\nAlthough the actual visual separation is not as clear as for JMDS,\nthe classification accuracy is increased.\nTo improve the visual embeddings,\nwe now equip the feature spaces \nwith distances from\nthe unsupervised SCOT distance estimation routine \n(Demetci, et al, doi:10.1089/cmb.2021.0446).\nThe updated distances lead to \nthe following [new version](https://imgur.com/a/U1E7GL7) of Figure 6.\n\nWe hope that we have addressed your questions.", "claims_and_evidence": "The proofs of Propositions 3.1-3.5 related to the formulation and solution of EW_lambda are clear, and the block-coordinate descent method based on Sinkhorn is reasonable. However, I think there is a lack of comparison with the GW distance, including aspects such as computational complexity, runtime, convergence, and the alignment effect.", "methods_and_evaluation_criteria": "The experiments on joint embedding of 3D shapes, alignment of feature spaces, and alignment of Gaussian mixture models are reasonable and can verify the capability of joint embeddings in both Euclidean and non-Euclidean spaces.  However, I'm curious as to why there was no comparison made with the alignment based on the GW distance.", "theoretical_claims": "I have reviewed the proofs of Propositions 3.1-3.5 and implementation of Algorithm 1, and I think they are reasonable.", "experimental_designs_or_analyses": "The experiment is reasonable, and the effectiveness of EW_lambda can be seen in various alignment experiments. Is it necessary to have a comparison with the baselines based on GW distance?", "relation_to_broader_scientific_literature": "Embedding different domains into a common space is beneficial for comparing different features, and it has practical application value. The  proposed method seeks near-isometric joint embeddings of the inputs into a metric space, which relaxes the constraints compared with Sturm's GW distance.", "essential_references_not_discussed": "I think the relevant literature has been covered.", "other_strengths_and_weaknesses": "Pros:\n\nThe paper presents the EW_lambda method which can map data from two distinct domains without known correspondences into a common metric space. It also theoretically proves the existence of a minimizer for this problem. Moreover, a bi-convex relaxation based on Sinkhorn is employed to numerically solve the issue. Additionally, the relationship with JMDS is discussed and the effectiveness of the method is experimentally validated.\n\nCons:\n\n(1) There is a lack of in-depth discussion on how to select the lambda parameter according to different datasets/features. The sensitivity analysis of the regularization parameter should be analyzed.\n\n(2) There are uncertainties about whether there are limitations when this method is extended to large datasets or high-dimensional features.\n\n(3) There is a comparison with GW in the method section, but there may be a lack of GW-based alignment baselines in \"Numerical Results\" section.\n\n(4) The bi-convex relaxation and block coordinate descent methods are based on the Sinkhorn algorithm. Are there any cases of numerical instability for this algorithm in practical applications? If so, how can these issues be addressed?\n\n(5) From the results of Figure 6, compared to JMDS, the proposed method fails to effectively represent the clustering structure of the dataset.", "other_comments_or_suggestions": "There are no other suggestions here.", "questions_for_authors": "All the questions are in \"Other Strengths And Weaknesses\".", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fYVBvmNHtp", "reviewer": "Reviewer_FyDV", "summary": "This paper proposes a framework for jointly embedding and aligning two metric measure spaces into a common metric space. The authors build their formulation from well-known prior works in this area, namely the Sturms and Memolis variants of Gromov Wasserstein distance. The main contribution of this paper is focused on a relaxation that can handle arbitrary metric measure spaces (i.e., not necessarily isomorphic) resulting in a relaxed Embedded Wasserstein Distance $EW_\\lambda$. This ultimately results in an unbalanced optimal transport problem with Gromov–Wasserstein marginal penalization. The authors provide a derivation of a numerical approach to solving it using bi-convex relaxation. \n\nThe experiments are mostly synthetic and focus mainly on comparing with a conceptual counterpart - JMDS Chen et.al 2023. Specifically, the results are shown for joint embeddings of (1.) 3D shapes - FAUST, Swiss Roll, spherical and toroidal embeddings, (2.) Feature spaces and synthetic GMMs.\n\nThe key difference seems to be that JMDS applies a free support approach for obtaining the coordinates in the joint embedding whereas the proposed $EW_\\lambda$ opts for a fixed support instead and optimizing for the weights. Multiple proof of concept demonstrations validate this aspect and show aligned, meaningful joint embeddings of the metric spaces", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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'\nand for your constructive feedback!\nWe revised our manuscript according to your suggestions.\nPlease find our answers below.\n\n**Experiments:**\nWe agree that\nthe shown experiments are mainly proof-of-concepts.\nTo improve the numerical part,\nwe extend Section 6.2 by a [CPU runtime comparison](https://imgur.com/a/RwHMZXm)\nwith respect to the different methods.\nMoreover,\nwe extend the supplement by a\nnumerical sensitivity analysis\nwith respect to the relaxation ([$\\lambda$-Experiment](https://imgur.com/a/g1Thyrg)) \nand regularization parameters\n([$\\varepsilon$-Experiment](https://imgur.com/a/odRSmkl))\nas well as\nwith respect to the discretization of the joint space\n([$Z$-Experiment](https://imgur.com/a/rlnnkVr)).\nPlease also note the domain adaptation experiment in Appendix C.\nThanks for pointing out possible future applications\nwhich, at the moment, are beyond the scope of this paper.\n\n\n**Complexity:** \nIn comparison to GW,\nthe complexity of our method is increased\nsince we rely on a multi-marginal OT problem.\nHowever,\nplease note that\nGW gives only a correspondence between the measures,\nbut no direct information about a joint embedding.\nAs mentioned above,\nwe extend the feature space embedding experiment \nto compare the [runtime](https://imgur.com/a/RwHMZXm)\n between the proposed approach\nand other methods in the literature.\nMoreover,\nwe now study the runtime in dependence of the discretization of $Z$\nfor the 2d embeddings in Section 6.1 ([$Z$-Experiment](https://imgur.com/a/rlnnkVr)).\n\n**Structure:**\n- Comment on page 1:\n  With the lack of invariance \n  (with respect to shapes),\n  we mean that the computed Wasserstein plans\n  depend on the actual embeddings of the shapes.\n  In contrast,\n  the Gromov-Wasserstein plans are always \n  independent of the embeddings or representations\n  of the considered shapes.\n  In the revised version the comment is accordingly clarified.\n- Figure 3: \n  The figure illustrates the multi-marginal plan $\\alpha$,\n  which is, in fact, four-dimensional.\n  The labeled marginals $\\gamma_1$, $\\gamma_2$, and $\\pi$\n  give the relation to the optimization problem\n  behind the original EW formulation.\n  The colors match the coloring of Fig 2,\n  which means that\n  everything related to $\\mathbb{X}_1$ is red and\n  to $\\mathbb{X}_2$ blue.\n  We add this in the captures of the figures.\n- Lastly, we will look into improving the structure and writing for the final manuscript.\n\n**Questions:**\nThe reason why our embeddings are not centered is that\nwe show the push-forward to the complete chosen grid $Z$.\nFor JMDS,\nwe obtain two point clouds\nand manually select a region containing both,\nwhich leads to a 'centered' visualization.\nIn fact,\nour and JMDS embeddings are invariant to (joint) translations.\n\nWe hope that we have addressed your concerns.", "claims_and_evidence": "Yes. I think the paper is theoretically constructive in its structure and new definitions are usually followed by a derivation of essential properties.", "methods_and_evaluation_criteria": "I think the experiments are comprehensive and a good proof of concept with multiple figures showing good joint embeddings. However, I do miss a broader applicability of the approach with respect to say applications like non-rigid 3D shape matching, transfer learning, domain adaptation, etc.", "theoretical_claims": "The important ones like the definition and properties of $EW_\\lambda$", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "To the best of my knowledge - I think the paper has been positioned quite well. The remarks on Lines 295 - 307 on Page 6 were the most refreshing: clearly outlining the key differences with JMDS. All things considered, I will not fault this submission too much for lack of novelty.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "This paper does feel lacking in a couple of aspects \n\n- The writing and structure are not optimal and somewhat unpolished. e.g. I am not sure I understand and agree with the comment on Page 1 \"However, classic OT lacks invariance to important transformations like rotations in Euclidean space\", \n- Line 297, typoJMDS and not JMSD,\n- I found Figure 3 very difficult to parse. What does it convey and why is the specific coloring? \n\n- Little discussion on the complexity of the method in comparison to other GW-based distances", "other_comments_or_suggestions": "Overall, I found this to be an interesting submission. It is theoretically comprehensive and conceptually enriches many prior OT-based distances. However, I won't categorize it with a clear acceptance - since it could be revised structurally for much more impact. Furthermore,  the experiments are much too synthetic and without any \"outstanding\" demonstrations clearly showing the real benefit of $EW_\\lambda$ over others. Also, I am missing any computational complexity issues like runtime, memory, etc.", "questions_for_authors": "Figures 4 and 5 why are the embeddings of the proposed method not centered unlike JMDS?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "H27yu6lQbe", "reviewer": "Reviewer_haA5", "summary": "The authors formulate an unbalanced OT problem where the marginals are penalized with GW distances and prove that their function have the minimizer. The authors show that if the regularization parameter goes to infinity, the solution their optimization problem is the embedded Wasserstein distance. To approximate the solution to their unbalanced OT problem, they define a relaxed bi-convex relaxation and present an algorithm using block-coordinate descent and multi-marginal entropy regularized OT. As their formulation seeks near-isometric joint embeddings into a metric space, the authors can use their method to find joint embeddings of data with different feature spaces. They utilize their method for some artificial 3D shape datasets and to align single cell sequencing data.\n\n### Update after rebuttal\n\nThanks to the authors for their responses. I keep my score", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for the positive evaluation \nand your constructive feedback!\nWe have revised our manuscript according to your suggestions in particular\nwith respect to quantitative experiments. \nPlease find our detailed answers to your questions below.\n\n**Quantitative FAUST:**\nWe agree that \nthe results shown in Figure 4 and 5\nare qualitative\nand can only give fleeting impressions\nof the algorithms. \nSince our method and JMDS can both be derived \nfrom the same functional,\nwe compare the achieved objectives,\nthe Wasserstein distances between embeddings,\nand the GW distances between embeddings and input spaces.\nThe experiment is repeated \nfor each pairwise combination \nof the first ten shapes from the FAUST dataset\nand different $\\lambda$.\nFor comparison, \nthe shapes are downsampled to 35 vertices,\nand both methods are regularized with $\\varepsilon=0.001$.\nFor our method, we set $Z$ according to a uniform\n30$\\times$30 grid on $[0, 1.3]$.\nThe average and standard deviation of\nthe GW distance (input space to embedding) \nand Wasserstein distance (between embeddings)\nare as follows: [Quantitative FAUST](https://imgur.com/a/x3EVFwa).\nIn total,\nwe here achieve better alignments than JMDS both in terms of GW and Wasserstein. \nWhile our method requires 4-5 seconds per embedding, JMDS requires 2-3 seconds.\nThese results \nare added to the revised version.\n\n**Separability:**\nWe agree that \nthe classes in Figure 6 are not clearly separated.\nNevertheless,\nthe constructed 2d embeddings increases the KNN accuracy \ncompared with the other methods.\nThe message behind Figure 6 is that\nthe embeddings of the two datasets \nfrom completely different feature spaces\nlooks quite similar\nand that the classes are mapped to common regions,\ni.e.,\nthe embeddings have similar color codings.\nTo improve the visual embeddings,\nwe now equip the feature spaces \nwith distances from\nthe unsupervised SCOT distance estimation routine \n(Demetci, et al, doi:10.1089/cmb.2021.0446).\nThe updated distances lead to \nthe following [new version](https://imgur.com/a/U1E7GL7) of Figure 6.\n\n\n**Experiments:**\nPlease note that there are additional experiments in the supplement\n(comparison between GW, $\\textrm{EW}_\\lambda$ and Wasserstein in Appendix B;\ndomain adaption experiment in Appendix C).\nTo further improve the numerical section,\nwe add the above quantitative study of the FAUST dataset. \nMoreover, we add illustrative\nsupplementary sensitivity studies\nwith respect to \nthe relaxation parameter ([$\\lambda$-Experiment](https://imgur.com/a/g1Thyrg)),\nthe entropic parameter ([$\\varepsilon$-Experiment](https://imgur.com/a/odRSmkl))\nand the discretization of the joint space  ([$Z$-Experiment](https://imgur.com/a/rlnnkVr)).\n\n\n**Minor Comments:**\n- Please notice that the push-forward is introduced \n  in the first paragraph of Section 2.\n  To improve the presentation,\n  we set the definition in\n  its own equation.\n- The grammatical error is corrected now.\n  \n**Questions:**\n- In our experiments,\n  Alg 1 needs less than 50 iterations to converge.\n  In the revised version,\n  we show the convergence of the objective over iterations\n  to our $\\varepsilon$ sensitivity study\n  ([$\\varepsilon$-Experiment](https://imgur.com/a/odRSmkl)).\n- Unfortunately,\n  there is so far no theoretical guarantee \n  about the convergence of Alg 1,\n  which we hope to find in our future work. \n  \nWe hope that we met all your concerns in the revised manuscript.", "claims_and_evidence": "The claims in the submission are supported.", "methods_and_evaluation_criteria": "The evaluation criteria make sense for the problem at hand.", "theoretical_claims": "I did not check the proofs.", "experimental_designs_or_analyses": "I checked the experiments. I am not quite sure what to make of the results in section 6.1. It seems like the results are in summarized Figure 4 and Figure 5 and they are more qualitative. Is there a way to compare the performance of JMDS against the author's method on the FAUST dataset across a larger set of pairs of human poses? Are the results in Figure 5 true across all pairs of poses in FAUST? Additionally, while Table 1 shows some convincing performance gains for their method over JMDS, I am not sure what the message of Figure 6 is? I do not see a very clean separation of cell types in either JMDS or the author's method.", "relation_to_broader_scientific_literature": "This paper seems to be most closely related to the joint multi-dimensional scaling algorithm which also uses OT to produce embeddings of different datasets into a shared Euclidean space. The authors provide a new framework for joint metric embeddings by using \\emph{unbalanced} OT and Gromov-Wasserstein distance (to enforce the embedding space to be nearly isometric to the inputs).", "essential_references_not_discussed": "I do not know of any essential references not discussed.", "other_strengths_and_weaknesses": "Strengths: The framework that the authors present seems to be novel as they produce their metric embeddings by defining a unbalanced OT problem with Gromov-Wasserstein penalization, which has not been previously done. They also connect their work to previous work as they show their optimization problem is the same as a previously defined \"embedding Wasserstein metric\" as the regularization term goes to infinity. \n\nWeaknesses: The experiments seem to be a bit thin, as I mentioned above.", "other_comments_or_suggestions": "1. The authors consistently use the notation $I_{1, \\sharp} \\xi_1$ which I believe is the push-forward measure but I think this notation is probably not familiar to everyone in the ICML audience so the authors might want to make a note. \n\n2. Grammatical error on line 136: \"the following relaxation enables to handle arbitrary....\" There should be another word between \"enables\" and \"to\".", "questions_for_authors": "1. Do we know how quickly algorithm 1 converges? \n\n2. Are there any guarantees about the quality of the approximation that Algorithm 1 produces?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OFNBdPBkjj", "reviewer": "Reviewer_CPxj", "summary": "In this paper the authors propose a novel alignment procedure  procedure across different spaces by optimizing a Wasserstein distance in a joint embedding space between two distributions that are close  in the GW sens to two distributions.  The idea is similar to the Sturm variant of GW and the authors express the problem as the optimization of the weighs of the distributions on a fixed support (calling it GW marginal penalization). They reformulate the problem as a multi-marginal OT problem and  propose a block coordinate descent algorithm to solve the optimization using at each step unbalanced MMOT solvers. A nice connection between the proposed approach and the existing Joint-MDS is discussed (joint MDS estimated support in euclidean space while the proposed approach estimate weights at fixed support). Several proof of concept of the approach illustrate what can be done with it in practice with notably very good performances in single cell alignment.\n\n## update after rebuttal\n\nThanks to the authors for their rebuttal. It clarified a few questions and the new results are appreciated but some proposed changes in the paper cannot be properly reveiwed so I will keep my score.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for your positive report and the constructive comments.\nWe revised our manuscript accordingly,\nespecially with respect to the limiting cases\nand additional references.\nPlease find our detailed answers below.\n\n**Limiting Cases:**\nThanks for pointing this out,\nwhich extends the limiting case ($\\lambda\\to\\infty$)\nfor existing isometric embeddings\ndiscussed in Prop 3.4.\nAs you suggested,\nthe limit ($\\lambda\\to0$) actually yields\na fixed support barycenter,\nand the limit ($\\lambda\\to\\infty$ without isometries)\ngives\nthe best possible projection\nfor both spaces separately.\nIn the revised version,\nwe extend Section 3\nby adding a new proposition.\nThe rigorous proof\nrelies on a subsequence argument\nand the weak continuity \nof the EW functional\nsimilar to the proof of Prop 3.4.\nThe proof is added to the supplement.\n  \n**Extendability:**\nFrom a theoretical point of view,\nthe proposed variational formulation can be adapted\nto handle multiple input spaces.\nThe main issue hereby is\nhow to compare the different projections on $Z$. \nIntuitively,\nwe may rely on pairwise Wasserstein distances;\nbut this approach does not yield\na 'tree-like' cost function \nfor the multi-marginal reformulation.\nAs a consequence,\nthe employed multi-marginal Sinkhorn method\n(Beier et al, 2022, doi:10.1007/s10851-022-01126-7)\ncannot be applied.\nAn alternative would be to utilize\nonly the Wasserstein distance \nbetween the $k$th and $(k+1)$st projection.\nSince the employed multi-marginal Sinkhorn method\nlinearly scales in the number of spaces,\nwe expect that the resulting method\nremains numerically tractable\nfor small numbers of spaces.\nThe actual choice of the model\nand the numerical implementation\nof a multiple space variant,\nhowever,\nrequires further studies.\nThanks for pointing out this imported extension,\nwhich is discussed \nin the conclusion\nof the revised manuscript.\n\n**Multi-Marginal Reformulation:**\nIn the final version,\nthe rewriting of the original problem is extended\nto give a better intuition behind the equivalence.\nThe detailed proof of Prop 3.5 is added \nto the supplement.\n\n**Choice of Geometry**\nYou are right: \nthe geometry of $Z$ has to be chosen\nwith respect to the problem at hand.\nTo study the sensitivity \nwith respect to the discretization of $Z$,\nwe repeated the experiments in Section 6.1\n(unrolling, 2d embedding)\nwith different grid sizes.\nThe results remain stable\nmeaning that we obtain the reported embeddings\nin different resolutions\n([$Z$-Experiment](https://imgur.com/a/rlnnkVr)).\nThe chosen grid, \nhowever,\nhas a major impact on the computation time.\nThese experiments are added to the supplement.\nPlease notice\nthat,\nbesides the alignments on the sphere and torus,\nalso the alignment of the Gaussian mixture models \nrely on a non-Euclidean geometry\nsince the alignment is calculated\nin the space of Gaussians\nequipped with the Wasserstein distance itself.\nWe agree that\nthese experiments are only proof-of-concepts\nthat can only point in the direction of possible applications\nand that real-world applications require further study.\n\n**Parameter and Regularization:**\nTo study the impact of $\\lambda$ and the entropic regularization,\nwe extend the experiments in section 6.1.\nAs you pointed out,\nsmall $\\lambda$ leads to fixed support barycenters\nand large $\\lambda$ to GW-minimizing projections\nof the input spaces\n([$\\lambda$-Experiment](https://imgur.com/a/g1Thyrg)),\nwhich nicely relates to the theoretical limiting cases.\nTo study the impact of the regularization,\nwe again adapt the experiments in Section 6.1\n([$\\varepsilon$-Experiment](https://imgur.com/a/odRSmkl)).\nThe results are added to the supplement, \nand we discuss the impact of entropic regularization.\n\n**Code and Reproducibility:**\nThanks for pointing out \nthat the parameter ranges \nfor the grid searches \nare missing.\nWe use the same distance graph and the same entropic parameter \n$\\varepsilon=0.001$ for all algorithms and four hyperparameter configurations per method. \nWe test $\\lambda=0.01,0.1,10,100$ for JMDS and our method. For UnionCOM, we use combinations of $\\beta=1,10$ and \na perplexity of 30 and 100 for the t-SNE dimensionality reduction. \nSCOT only relies on $\\varepsilon$ and the input distances and has no additional hyperparameters.\nIn the revised version,\nthe ranges are added.\nAlso,\nour implementations will be publicly available via Github\nto reproduce or reuse our implementations.\nDue to the anonymity requirement and \na short e-mail exchange on the official guidelines,\nwe cannot link the code here,\nbut the link will be provided in the final version.\n\n**References:**\nThanks for pointing out the related references\nwhich we incorporate in the revised version.\nRegarding the GW dimension reduction,\nwe also refer to (Clark, Needham, Weighill: arXiv:2405.15959v2),\nwhere the relation between GW and MDS is discussed.\n\nWe would like to thank you again for your suggestions\nthat significantly improved the manuscript\nand hope that you can accept the paper now.", "claims_and_evidence": "The claims in the abstract are all reasonable and well supported. The empirical evidence about the interest of the method is a bit limited since as stated it relies a lot on several proof-of concept but they illustrate well the method and the diversity of potential application.", "methods_and_evaluation_criteria": "As stated above the numerical experiments are proof of concept and rely  mostly on qualitative visualizations expect for the application to single cell alignment where classical performance measures for the field are reported (and the proposed approach use a very specific approach using barycenter mapping.)", "theoretical_claims": "The theoretical claim seem to be correct from a quick pass over the proofs in appendix but I did not check them in detail. The transition from the original problem to a multi-marginal problem is a bit quick and the proof of proposition 3.5 that is important in practice is missing so the supplementary would be improved with those details (that are not as obvious as stated in the paper).", "experimental_designs_or_analyses": "The experiments are interesting. The visualizations are pertinent and show the interest of the method. On the other hand by doing multiple proofs of concept, the authors limit their study and barely discuss or illustrate the impact of some very important parameter such as the effect of lambda and of the choice of the fixed support Z_i.\n\nAlso the results on the single cell alignment are interesting but not really reproducible without the code since no parameter validation grids are provided.", "relation_to_broader_scientific_literature": "The paper is obviously strongly related to recent extension  of MDS that is very well discussed and positioned. But some references from the OT community are missing and would deserved to be discussed.", "essential_references_not_discussed": "An important reference that should be discussed because the proposition of estimating the weighs minimizing GW on a fixed support is illustrated in their Figure 2: The proposed approach is a clear extension that handle multiple distributions instead of one but the connection is clear and more general metrics on Z are proposed in the reference\n\nVan Assel, Hugues, et al. \"Distributional Reduction: Unifying Dimensionality Reduction and Clustering with Gromov-Wasserstein.\" arXiv preprint arXiv:2402.02239 (2024).\n\nThe low rank OT methods use a similar final factorization between the OT plans (where the rank r is the size oif the support in Z for the proposed method). For instance Forrow can be seen as a Wasserstein version of the joint MDS.\n\nScetbon, Meyer, Gabriel Peyré, and Marco Cuturi. \"Linear-time gromov wasserstein distances using low rank couplings and costs.\" International Conference on Machine Learning. PMLR, 2022.\n\nForrow, Aden, et al. \"Statistical optimal transport via factored couplings.\" The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019.\n\nHalmos, Peter, et al. \"Low-rank optimal transport through factor relaxation with latent coupling.\" Advances in Neural Information Processing Systems 37 (2024): 114374-114433.\n\nA very similar the MMOT formulation with  specific structure was proposed in :\nKerdoncuff, Tanguy, et al. \"Optimal tensor transport.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 36. No. 7. 2022.", "other_strengths_and_weaknesses": "+ The method is interesting, the \"projection\" onto a joint space and the use of Wasserstein is an intriguing way to transfer information between different source distribution in different spaces and to recover a global alignment through the shared space. \n\n+ The limit case of the proposed approach should be discussed more in detail. if lambda goes to infinity then the solution and the most isometric preserving projection that minimize the Wass distance (but the two projections P_#\\pi can be very different). When lambda goes to 0 then the problem seems to revert to a GW barycenter with fixed support since the Wasserstein must go to 0 so you search for a unique distribution that is close to the two in the GW sens hence a barycenter. Those special case are important to discuss (an prove maybe results?) so that one can have a better idea of the positioning wrt the very active OT GW extension community.\n\n+ How extendable is the proposed approach to multiple distributions? The method is already quite heavy with a very high dimensional MMOT and does not seem to scale. This is not a problem but it needs to be discussed.\n\n+ The proposed reformulation an optimization algorithm are also interesting. But more details about why there is an equivalence between the MMOT formulation and the factorized OT plan is important (due to the structure of the global objective value but not as obvious as stated).\n\n+ The choice of the geometry and support in Z is very sensitive. The experiments are carefully selected so that either the geometry corresponds to the distribution (unroll in 2D, torus, 2D euclidean known to work well for single cell). But outside of the very toy example of sub-part of torus, the geometry is barely used in practice or remains euclidean so the argument that the proposed method is more general because it can be used for non euclidean geometry is not well justified in practice. Especially since learning the geometry can be done implicitly with estimation a GW barycenter that provides pairwise \"distances\" on the learned manifold and also allows projection on a smaller dimension in euclidean.\n\n+ Numerical experiments are mostly proof of concepts. They are interesting and clearly the proposed strategy of fixing the support can be of interest wrt JMDS. But on the other hand it is a bit difficult to estimate the impact and practical use of the proposed method outside of single cell alignment (and even in this case it is compared only to few methods). This limits a little the potential impact of the method.\n\n+ As in many OT papers the authors do a complete modeling of the problem they want to solve and at the end just add some entropy regularization when numerical optimization becomes necessary. This behavior should not be promoted since entropic regularization has an important impact on the final OT plans and a study of the impact of this regularization in terms of mass spread/performance VS computational time is necessary to evaluate its impact. As noted above this kind of parameter sensitivity analysis (also wrt lambda and the support of Z) is important for the practitioner.", "other_comments_or_suggestions": "See above", "questions_for_authors": "See above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In this submission, the authors proposed a new OT-based method to embed and align incomparable geometric objects jointly, leading to an unbalanced OT (UOT) problem with a GW-based marginal penalty. The proposed method leads to an optimization problem solved by block coordinate descent. Experimental results show the potential of the proposed method for aligning Euclidean and non-Euclidean objects.

The scores of the proposed method are mixed, but most reviewers appreciate the idea and admit the novelty of the method. However, AC and some reviewers have concerns about the computational efficiency of the method, especially when dealing with large-scale and high-dimensional objects. These concerns are not fully resolved in the rebuttal phase.

In summary, AC recommends to "weakly accept" this work.
Decision: Accept (poster)