Disentangled Representation Learning with the Gromov-Monge Gap
Authors: Théo Uscidda
Luca Eyring
Karsten Roth
Fabian J Theis
Zeynep Akata
marco cuturi
Venue: ICLR 2025 Poster (2025)
Keywords: Disentangled Representational Learning
Optimal Transport
Geometric Learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We develop a disentangled representation learning method that efficiently combines prior matching with geometric constraints on the latent space, framed through the lens of quadratic optimal transport.

Abstract: Learning disentangled representations from unlabelled data is a fundamental challenge in machine learning. Solving it may unlock other problems, such as generalization, interpretability, or fairness. Although remarkably challenging to solve in theory, disentanglement is often achieved in practice through prior matching. Furthermore, recent works have shown that prior matching approaches can be enhanced by leveraging geometrical considerations, e.g., by learning representations that preserve geometric features of the data, such as distances or angles between points. However, matching the prior while preserving geometric features is challenging, as a mapping that *fully* preserves these features while aligning the data distribution with the prior does not exist in general. To address these challenges, we introduce a novel approach to disentangled representation learning based on quadratic optimal transport. We formulate the problem using Gromov-Monge maps that transport one distribution onto another with minimal distortion of predefined geometric features, preserving them *as much as can be achieved*. To compute such maps, we propose the Gromov-Monge-Gap (GMG), a regularizer quantifying whether a map moves a reference distribution with minimal geometry distortion. We demonstrate the effectiveness of our approach for disentanglement across four standard benchmarks, outperforming other methods leveraging geometric considerations.

Published as a conference paper at ICLR 2025
DISENTANGLED REPRESENTATION LEARNING WITH THE GROMOV-MONGE GAP
Theo Uscidda 1.2
Luca Syyring.3.44. Karsten Roth 2.1,5
Fabian Theis 2,3,4
Zeynep 1.34.
Marco Cuturi 1,6
*CREST-ENSAE
"Helmholtz Munich
TT Munich
"Munich Center of Machine Learning thec, use dededddaeeeeeeeeeeee fr
TTubingen Al Center
6 Apple
luca. moca.eyring@rrr eyring@tum.de
ABSTRACT
Learning disentangled representations from unlabelled data is fundamental chal- lenge in machine learning. Solving it may unlock other problems, such as general- ization, interpretability,
fairness. Although remarkably challenging solve in theory, disentanglementt is often achieved in practice through prior matching. Fur- thermore, recent works have shown that prior matching approaches can be enhanced by leveraging geometrical considerations, c.g., by learning representations that preserve geometric features of the data, such as distances or angles between points. However, matching the prior while preserving geometric features is challenging, a mapping that fully preserves these features while aligning the data distribution with the prior does not exist in general. To address these challenges, novel approach to disentangled representation learning based on quadratic optimal transport. We formulate the problem using Gromov-Monge maps that transport distribution onto another with minimal distortion of predefined geometric features, preserving them as much as can be achieved. To compute such maps, we propose the Promov-Monge-Giag (GMG), regularizer quantifying whether map moves a reference distribution with minimal geometry distortion. We demonstrate the effectiveness of our approach for disentanglement across four standard benchmarks, outperforming other methods leveraging geometric considerations.
introduce a
INTRODUCTION
Learning low-dimensional representations of high-dimensional data is a fundamental challenge in unsupervised deep learning (Bengio et al., 2014). Emphasis is put on learning representations
that allow for efficient and robust adaptation across a wide range of tasks (Higgins et al., 2018
Locatello et al.] 2019ah. The fundamental property of disentanglement has shown significant promise
to improve generalization (Locatello et al.) 2020 Roth et al. 20233: Hsu et al.
2023
Barin-Pacela
al. 2024), interpretability and fairness (Locatello et al. 2019b] Träuble et al.
2021
Most works
regard disentanglement as a one-to-one map between learned representations and ground-truth latent
factors, effectively seeking to recover these factors from data alone an unsupervised fashion. While
unsupervised disentangiement is theoretically impossible (Locatello et al. (2019)), the inductive biases
of autoencoder architectures ensure effective disentanglement practice (Rolinek al. 2019 Zietlow et al. 20221. Most approaches operate using variational autoencoder (VAE) frameworks Kingma and Welling 2014), using objectives that match latent VAE posteriors to factorized priors (Higgins
et al. 2017 Kim and Mnih 2018 Kumar et al. 2018 Burgess et al. 2018 Chen et al. 2018).
More recently, studies such as Gropp et al. (2020): Chen et al. (2020a): Lee et al. (2022): Horan et al.] (2021); Nakagawa et al. (2023); ;; Huh al. (2023); ); Hahm et al. (2024) have provided new perspective, showing that geometric constraints on representation spaces may also enable disentanglement. Typically, latent representations
encouraged to preserve key geometric features of the data distribution, such as (scaled) distances or angles between samples. Horan et al.) (2021) even demonstrate that unsupervised disentanglement is always possible provided that the latent space
"qqual contribution
I Equal advising

---

Published as a conference paper at ICLR 2025
is locally isometric to the data, further supporting the geometric desiderata. However, combining prior matching with these geometric aspects is challenging. In general. a mapping that perfectly aligns the data distribution with the prior while fully preserving the geometric features of interest may not exist. This leads to an inherent trade-off:: Practitioners must carefully fine-tune regularization terms, either by altering prior matching to prioritize geometry preservation, vice-versa.
In this work, we demonstrate how to effectively combine geometric desiderata with prior matching within the VAE framework, using optimal transport (OT) theory (Santambrogio 2015 Peyré and Cuturi 2019). By treating mappings from the data manifold to the latent space (encoders) or vice versa (decoders) as transport maps T X + +V, we can leverage the Gromov- (Monge. Wasserstein] paradigm (Stur 20233 Mémoli, (011). which aligns two distributions by finding mapping that minimizes the distortion between intra-domain cost functions defined on their supports. Specifically, we consider cost functions cx cx(x,x' ) on A' and yy) on )' that encode geometric features such as scaled distances or angles. Consequently, the resulting mapping transforms one distribution onto the other while preserving these geometrie features much as possible.
Our Contribution: A novel OT-based approach to disentanglement through geometric considerations.
(i) We address the challenge of learning disentangled representations using geometric constraints by leveraging Gromov-Monge mappings between the data and prior distributions. Since fully preserving geometric features- -uch as (scaled) distances or angles between points-during the alignment of these two distributions is generally impossible, we aim to find an alignment that, instead, minimizes the distortion of these features, thereby preserving them as much as possible.
(ii) Inspired by (Uscidda and Cuturi, 2023), we introduce the Gromov-Monge Gap (GMG), regularizer that measures how closely a map T approximates a Gromov-Monge map for costs cx: cy. GMG measures whether I" transports distributions with minimal distortion w.r.t. CX. cy. We propose an efficient procedure to compute GMG and describe how to integrate it within the VAE framework.
(iii) We show that when cx and ey encode scaled distances or angles, the GMG and its simite-sample counterpart weakly convex functions. In both cases, we precisely characterize the weak convexity constants and analyze their practical implications for practitioners.
(iv) Across four standard disentangled representation learning benchmarks, we show that incorporating desiderata via the GMG significantly enhances disentanglement across various methods, from the standard 8-VAE to the combination of 3-TCVAE with HFS (Roth et al. 2023).
BACKGROUND: ON DISENTANGLEMENT, QUADRATIC-OT AND DISTORTION
2.1 DISENTANGLED REPRESENTATION LEARNING
The Disentanglement Formalism. Disentanglement has varying operational definitions. In this work, we follow the common understanding (Locatello et al. 2019a] 2020 Roth et al. 2023] (Träuble et al. 2021 Higgins et al. 2017) where data is generated by a process p x (z) operating ground- truth latent factors Z ~ pz). modeling underlying source of variations (s.a. object shape, color, background.... )). Given dataset D { X, Pastas unsupervised disentangled representation learning aims to find a mapping es s.t. e (x;) R E[|[|x;] up to element-wise transformations. Notably, this is to be achieved without prior information on p(z) and p(x(z).
Unsupervised Disentanglement through Prior Matching. Most unsupervised disentanglement methods operate on variational autoencoders (VAEs, Kingma and Welling (0144), which define a gen- erative model of the form (x z) p(z) (xlz). Here, pa (x z)) is a product of exponential family distributions with parameters computed by a decoder de (z). The latent prior p(z) is usually chosen as standard Gaussian A (z)00,, Ia), and the probabilistic encoder Ts (z x) is implemented through neu- ral networks es (x), (x) that predicts the latent parameters so that qo (z|x) (Z(ee, o (x)). The 3- VAE (Higgins et al. 2017) achieves disentanglement by minimizing
min E (x)) 8.0
log Po (x|z) + D Dg (z(xx [p(z))
(1)
(i) reconstruction
(ii) prior matching
which enforces w-wehtteed prior matching on top of the reconstruction loss, assuming statistical factor independence (Roth et al. 2023). Several follow-ups refine latent prior matching through
2

---

Published as a conference paper at ICLR 2025
different losses or prior choices (Rolinek et al.) 2019 Kim and Mnih 2018: Burgess et al. 2018 Chen et al. 2018 Moor et al. 2021 Balabin et al. 2024
Kumar et al. 2018
Disentanglement through Geometric Lens.
Recent studies have revealed a fundamental connec- tion between geometric structure preservation and disentanglement in learned representations (Gropp
et al. 2020; Chen et al. 2020a Lee et al. 2022, Nakagawa et al., 2023 Huh et al.. 2023). This
connection was theoretically established by Horan et (2021) proving that unsupervised disentangle-
ment is always feasible when the generative factors are sufficiently non-Gaussian and maintain local isometry to the data. Our work builds directly on this insight by developing a learning framework that promotes representations that are close as possible to being isometric to the data. To quantify geometric preservation between spaces of different dimensions, we leverage quadratic OT theory. which originated in Koopmans and Beckmann (1957) and was formalized by [Mémoli] (2011) as a framework for measuring isometric correspondence between metric spaces. We detail these tools in Section 2.2 before showing how they can be used to learn representations in Section 2.3 Additionally, in concurrent work, (Sotiropoulou and Alvarez-Melis (2024) also introduced the Gromov-Monge work.,Sotiropoule
gap. However, they do not apply it for disentangled representational learning.
2.2 QUADRATIC OPTIMAL TRANSPORT
theory studies efficient ways to map probability distribution onto another. Linear OT formulations, such as the Mongel (1811) problem, require domains X. that can be directly compared through cost function c(x: y) defined between their elements. When these distributions lie on incomparable domains, one must instead rely on quadratic formulations of OT (Q-OT), which instead compare geometric structure through intra-domain costs, also known as the Gromov-Monge (GM) and GW problems. In the context of representation learning, representation and data spaces are incomparable by design, which necessitates the use of Q-OT in this work.
OT Peyyé and Cuturi 2019)
Gromov-[Mongs Wasserstein] Formulations. Consider two compact X' C Rdx, y C Rdy, equipped with intra-domain cost cx : X x X IR and ey y y + R. We assume that cx and Cy (or -cx and -cy) are CPD kernels (Def. (A.I). For p E (x') and E P())), two distributions supported on each domain, the GM problem (Mmoli and Needham (002)) seeks a map T X y that fush-forwards p onto q. while minimizing the distortion of the costs:
inf / 1 |xxx(x cy T(x))) dp(x))((xx)
(GMP)
T-T:== JxxX
When it exists, we call a solution T to GMP)) a Gromov-Monge map for costs cx, cy. However, this formulation is ill-suited for discrete distributions p, q. the constraint set might be empty in that case. Replacing maps by coupling π E [p q). i.e. distributions on X' x y with marginals p and g.
we obtain the GW problem (Sturm. 2023 Mémoli,,200
GW g) == min {
1 [xxx(xx x') - xy)))')) di (x,
dr (x'xx')))
(GWP)
=TIl(pa)) (xxx))²
A solution nT to (GWP) always exists, making GW(p, q) a well-defined quantity. It quantifies the minimal distortion the geometries induced by ex and ey achievable when coupling p and q.
translates Discrete Solvers. When both p and & are instantiated as samples, GW Prob. transaals to
a quadratic assignment problem, whose objective can be regularized using entropy (Cuturi 2013) Peyré et al. 2016). For empirical measures pn. > 8x.s
= + トッ 8y, and 2 0, we set:
In
GWe(Pn: gn) == min PEUn
Σ
(Cx. Cy ) P. Prr cH(P),
(EGWP)
3'=1
with Cx = [xx xyy )is,, Cy = 11 E Ruxn Un = (P B PI = 1' 11,} and (P) - E. P. log(P..)) (.j). As O. recover GWexiey GW Entropic regularization improves computational performance, as can solve (EGWP) using scheme that iterates the Sinkhorn algorithm (see Appendix D.3 for full details). This solver has Onn2) memory complexity. Its time complexity depends on Cx, Cy. For general cx., cy. it runs in O(n))
However, for the most common practical choices cx = Gy = (.. ( or Cx Cy = - -. 12, it can also be reduced Onn (dx + dy:)). as detailed by Scethon et al. 2022 $3 & Alg. 2).
3

---

Published as a conference paper at ICLR 2025
2.3 DISTORTION IN REPRESENTATION LEARNING
Given an arbitrary map T : X' + V' we consider how it can be learned to preserve predefined geometric features. In VAE, T may represent either the encoder eg, which generates latent codes from the data, or the decoder dg, which reconstructs the data from these codes, In the case of the encoder, X' corresponds the data, and y is the latent space, while for the decoder, these roles are swapped. Assuming that do perfectly reconstructs the data from the latents produced by ea. i.c., es dg Td., the preservation of geometric features by either cg or dg becomes equivalent. Therefore, in the following sections, we refer to T as either the encoder or the decoder without loss of generality,
We encode geometric features using a cost function for each domain: Cx : X x X - R and ey y x y + R. Ideally. T should preserve geometry, which means that T preserves costs, that is, cx x') R ey (T(x), Tx')) for x, x' E X. practice, two types of costs are often used:
[i] (Scaled) squared L2 distance: cx(x, x') ||x
x' x'122 and yy, a² y1l,, with a > O. A map T preserving Cx, cy preserves the scaled distances between the points, i.e. is scaled isometry. When = 1, we recover the standard definition of an isometry.
[ii] Cosine-Similariis
cx(xx,')
cos-sim(x, x')
< [x||a ||x'||2 and cy(y,y))
x
cos-sim(y, y') similarly. On has cos-sim(x, x') = cos(@x xox((x.x:) where 0x.x' is the angle between x and x'. A map I preserving cx, ey then preserves the angles between the points, i.c. it is conformal map. Note that if 7 is (scaled) isometry (see above), it is conformal map.
We refer to these costs via L2 for [i] with a = 1, ScL22 for [i] with a * and Cos for [ii]. Introducing reference distribution r E P(X), weighting the areas of X where we penalize deviations of cx cx x') from (Tx),,, T(x)), we can quantify this property using the following criterion:
Definition 2.1 (Distortion). The distortion (DST) of a map T, for cost functions ex, cy and reference distribution r, is defined as:
S Jxxx
DT), =
-
dr(x) dr(x').
(DST)
DST,(T) quantifies how much distorts geometric features encoded by cx, ey on the support r. that is, when (T) 0, one has cx(x, = (xx( for x, x' Spt(r).
Distortion as Loss for Representation Learning. Nakagawa et al. (023) suggest promoting geometry preservation by regularizing the encoder de using the DST. with ScL22 as costs, and the latent representation as reference distribution, namely r P While they use it within a WAE, we translate their objective to VAE setting adopted later in the paper. This results in:
min Ex Ex- 2x-paanz-00000le (x)00
log (x|z)) + 3DII (zx(<()))lll(x))) + ADST, (do)
A>0
(2)
(i) reconstruction
prior matching
(3i) geom. preservation
Given the choice of the costs, de is distortion-free (i.e., (dT, =
0), if it is a scaled isometry.
Challenges Arising from a Mixed Loss. Since a scaled isometry that maps the prior onto the data distribution may not exist, there is an inherent trade-off between minimizing terms (ii) and (iii), which are responsible for achieving practical disentanglement. As these terms cannot be simultaneously minimized to O, the DST loss will move away from accurately matching the prior, which will negatively impact the quality of the learned latent representations. This naturally raises the question Instead of seeking a distortion-fre decoder, we should seek
of how to avoid this over-penalizan decoder that transports the prior to the data distribution with minimal distortion of the costs Cx' Cy. In other words, the decoder should be Gromov-Monge map between the prior and the data distributions for costs cx, ey (see 2.2). Conversely, we choose to regularize the encoder, the same reasoning applies by swapping the roles of the prior and the data distribution. In the next section, introduce a regularizer to fit Gromov- -onge maps, which we will use as a replacement to the DST.
3 DISENTANGLEMENT WITH THE GROMOV-MONGE GAP
Building on these geometric preservation principles, we introduce in $3.1 the Cromov-Monge Gap (GMG), regularizer that measures whether map moves distributions while preserving geometric
4

---

Published as a conference paper at ICLR 2025
I'n
T'n
X2
X It?
X Rt?
T
(xx))
Tprn
T(xa)
7
T(x))
T(x)
y C
T(xa)
YCR
T(xa)
GMG,,, (T)
(a²
GMG,, (T) = DST (T)
2
+ +
(a
DST (T)
minimal DST
Figure 1: Sketch of GMG,, (T) for different maps T. We use discrete reference distribution on 3 points, and cx = Cy = ||. -- ||2. hence measuring if T minimally distorts the distances. On the left, T is the optimal map. I and maps the three points with minimal (yet non-zero) distortion, which is measured as the sum of the squared lengths of the orange segments. This results in GMG,,, (T) = On the right, T swaps two points compared to T specifically, X2 and xycausing higher distortion than the minimal one, and measured as the sum of the squared lengths of the red segments. This results in GMG,, (T)) > O, equal to the gap between these distortions.
features as much as possible, i.c., minimizing distortion while fitting the marginal constraints. $3.2 then shows how the GMG can be efficiently computed from samples be practically applicable in the VAE framework. This transitions into $3.3 studdyig (weak) convexity properties of the GMG, as an operator. Finally, in $3.4 we describe how integrate the GMG with disentangled representation learning objectives, effectively combining prior matching with geometric constraints.
3.1 AN EFFICIENT GAP FORMULATION FOR DISTORTION
introduced the Monge gap, regularizer that measures whether map T transports reference distribution the minimal displacement cost. Building this concept, we replace "displacement" with "distortion" to introduce the Gromov-Monge gap, regularizer that assesses whether a map T transports a reference distribution at the minimal distortion cost.
Recently, Uscidda and Cuturi 2023
Definition 3.1 (Gromov-Monn -Mnnee gap). The Gromov- Monoo gap (GMG) of map T, for cost functions cx,cy cy and reference distribution is defined as:
GMG, == DST, GW(r, Tpr)
(GMG)
We recall from Eq. (WPP) that GW(r.Tpr) represents the minimal distortion of cx, Cy achievable when transporting Tpr. Thus, the GMG quantifies the difference between the distortion incurred when transporting r to Tär via 7, and this minimal distortion. Formally, when Prob. (GMP) and Prob. (GWP) between r and Tér are equivalent, the the GMG is the suboptimality gap of T in Prob. (GMP) This is the case, for example, when r is a density and ex = ey = (., (Dumont et al. 2022). Otherwise, the GMG is the suboptimality gap of (Ia, fr in Prob. (GWP) between
and T . In light of this, it is a well-defined quantity and:
The GMG measures how close T is to be a Cromov-Monne map for costs Cx2, Cy. Indeed, GMG, 2 with equality if T' is Gromov-Monge map solution of Prob, (GMP) between and Tžr, i..., T' moves r with minimal (but eventually non zero) distortion; see App. B.
When transport without distortion is possible, the GMG coincides with the distortion. When there exists another map X ) transporting to Tír with zero distortion, i.c., Uir = Tpr and DST,(U) 0, then GMG, (7) DST, (T"). Indeed, GW(r, Tár) 0 in that case, the coupling π = (Td, U)tr sets the GW objective to zero, thereby minimizing it.
The last point illustrates how the GMG functions as a debiased distortion. It compares the distortion induced by T to a baseline distortion, defined the minimal achievable distortion when transforming r into Ttr. Thus, when transformation without distortion is achievable, this baseline becomes zero, and the GMG equals the distortion. Consequently, the GMG offers the optimal compromise: it avoids
5

---

Published as a conference paper at ICLR 2025
the werrpenalizztion induced by the distortion when fully preserving cx, cy is not feasible, yet it coincides with it when such full preservation is feasible. See Fig. I for a simple illustration.
The Influence of the Reference Distribution. A crucial property of DST, is that if T' transforms
without distortion, it will also apply distortion- to any distribution s whose support is contained within that of Formally, if DST, (T) = 0 and (()) with Spt(s) S Spt(r), then DST, (T)
0. This raises a question for the GMG: If T maps r with minimal distortion, does it similarly map with minimal distortion? We answer this question positively with Prop. (3.2).
Proposition 3.2. If (T) 0, Vs E P(X) s.t. Spt(s) C Spt(r), one has GMG, = 0.
3.2 ESTIMATION AND COMPUTATION FROM SAMPLES
Plug-In Estimation. In practice, we estimate Eq. DST and Eq. GGMG from samples X1,..., Xn ~ We consider the empirical version In 2i=1 0x. of and use plug-in estimators, i.e.
DST,, (T) R²
(cx(xx,x) ->)(((xx),,
(3)
5,=1
and GMG,, (T) = DST,, r.. (T) GW Ttra), where Tern = Eis ÓT(x))-
Efficient and Stable Computation. Computing the GMG requires solving discrete GW problem between Tn and Ter, to get GW (Tn,TTrrn) We com- pute this term using an entropic regu- larization & 2 0, as in Eq. (EGWP)
Algorithm 1 GMG(X).... Xp. T.E).
1: Require: XI,..., Xn ~ map T; entropic solver GW, entropic regularization scale Eo (default = 0.1), statistic operator cost matrix stat (default mean).
2: t,,...ty ty + T(x))........
(xx))
3: Cx + [xx(x,,x, (x)xx,x)))sssss Iss.ss
usually (n?
4: Cy loy(t,, ty Jissrss
usually (n"ddy)
GMGr. (T) == DST,. (T) (4)
5: DST + mean Cx Cy)3)
dn² usuully (n²) 7: GW <- GW(Cx. Cy.c) O(n) or (n + dy))
-GWe(rm; Tirn)
e + en stat Cx) stat
O, recover ( = GMG,, (T). More- over, the entropic estimator preserves positivity, as for 2 0, GMG, (T) 2 0 (see B.3 We compute GW (Fa, Ttrn) using Peyré et al. (2016)'s 's solver introduced in 2.2 We use the implementation provided by ot ott-jaa jat-jaxx (Cuturi et al. 2022). In practice, we select based on (positive) statistics from the cost matrices Cx,Cy, We define a scale Eu and set Eu- (Ca) ((Cx)-stt((C)) stat(Cy)). Standard options for the statistic include stat E {mean, max, st d]. This procedure is equivalent to running the entropic GW solver on the re-scaled cost matrices Cx Ate/statCC,, (xx) and Cy/stat (Cy) with E = En: see App. B.2 It is
Choosing
8: return DST GW
common practical trick, initially suggested for stabilizing the Sinkhorn algorithm Cuturi 2013).
Computational Complexity. For usual costs, including inner products, f distances, and standard CPD kernels, the computation of ex (x, x') xx) (resp. cy(y, y') can be done in O(dx) time (resp. (dyy)) Consequently, the DST can be computed in O(n" (dx + dy)) time. Furthermore, as discussed in & 2.2] the time complexity of the entropic GW solver is (n²)) in general, but can be reduced O (dd + dy)) when cx = Cy = (:5r), cx = || --. -12,cy
all.--.1l 11.- Therefore, since the cosine similarity is equivalent to the inner product, up to pre-normalization of x; and T(x;). this solver runs in Onn + dy)) for the costs of interest L2 and Cos. The complete algorithm, along with time complexity analysis of each step, is described in Alg. [ We stress that, in all cases and for any cost function, the time complexity depends linearly on the dimensions of the source and target spaces, dx and dy. making the GMG scalable to high- nigh-dimensional distributions.
3.3 (WEAK) CONVEXITY OF THE GROMOV-MONGE GAP
As laid out, the GMG can be used as a regularization loss to push any model to move distributions with minimal distortion. A natural question arises: is this regularizer convex? In the following, study the convexity of + GMG, (T) and its finite-sample counterpart T + GMG,, (T). We focus on the costs L2 and Cos. For simplicity, we replace Cos with (, +), these costs are equivalent, up

---

Published as a conference paper at ICLR 2025
to normalization of r and T. We respectively denote by GMG; and GMG the GMG for these costs. We start by introducing a weaker notion of convexity, previously defined on F (Davis et al. 2018), which we extend here to [2(r) == {T [ == = [ 12 dr(x) < +oo}.
Definition 3.3 (Weak convexity.). With 7 > functional F (2(r) R is y-weakly convex :T + [T (r) is convex.
weakly convex functional is convex up to an additive quadratic perturbation. The weak convexity constant y quantifies the magnitude of this perturbation and indicates degree of non-convexity of F.. A lower suggests that F is closer to being convex, while a higher 7 indicates greater non-convexity,
Theorem 3.4. Both GMG, and GMG as well as their finite sample versions, are weakly convex.
Finite sample. We note X E IRnd the matrix that stores the X i.e. the support of Tn. as rows. Then, (i) GMG and (ii) GMG are respectively (i) Y2.n and (ūi) "imerr, "immer.n-Weakll convex, [XX]] and 72.m Vinner, + max;=1...n ||x
where: "funer,n = Amax ( (xX Amin (
Asymptotic. (i) GMG,, and (ti) GMG are respectively (i) 7/2 and (ii) Wiaker convex, where: "(nner = Xmax (Exar [xx and Y2.n = Tiner + maXxESpt(r) [xx|2_
From a practitioner's perspective, we analyze the insights provided by Thm. 3.4 in three parts.
First, we have 7/2 2 Tinner. Therefore, GMG is less convex than GMG making it harder to optimize, and the same argument holds for their estimator. In other words, we provably recover that, in practice, preserving the (scaled) distances is harder than simply preserving the angles. Second, Vimner = Xmax(Exxx [xx I) 2 A (Cov [x]), this exhibits tradeoff w.r.t. Prop. (32) by choosing a bigger reference distribution т, we trade the convexity of the GMG. For 72, the dependency in even worse. In practice,
then choose r with support as small as possible, precisely where we want 7 to move points with minimal distortion.
Third, and most surprising, the finite sample GMG is more convex in high dimension. Indeed, Tinnern is the spectral width of 1 XXX" which contains the (rescaled) inner products between x
X; ~ r. When n > d, Amin (Xx = O as rank ) = d. Then, Tinner, increases, which in turn decreases the GMG's convexity. However, when
n, Amin (XXT)= > 0 if X is full rank. Intuitively, GMG is nearly convex when XXT is well conditioned. Assuming that the x; are normalized, this might happen in high dimension, as they will be orthogonal with high probability. This suggests that, contrary to the insights provided by the statistical OT literature (Weed and Bach 2017 Zhang al. 2023). the GMG might not benefit a large sample size.
3.4 LEARNING WITH THE GROMOV-MONGE GAP
General Learning Procedure. Given a source p and a target distribution q, we can use the GMG to guide a parameterized map To towards approximating Gromov-Monge map between p and g. We handle the marginal constraint Totp q separately through fitting loss AToo, p). Provided any reference s.t. Spt(p) c Spt(r), we minimize
min ATe, p. q) + XGMG,
(5)
A can operate on paired (e.g., in VAE, the reconstruction loss), or unpaired (e.g., in VAE, the KL. loss in VAE), samples of p and q. Note that, in theory and as stated in Prop. (3.2). we can select any reference such that Spt(p) c Spt(r). However, based the insights from Thm. (3.3), we typically choose r with minimal support size and, in practice, set r = p. This learning procedure is illustrated in Fig. 2 where also explore the effect of replacing (T), by DST,(T)) in Eq. (
VAE Learning Procedure. In the VAE setting, we can use the GMG promote the (i) encoder eg or the (ii) decoder do to mimic Gromov- Monge map. In (i) use r. Paata the data distribution as reference r, while in (ii) we use the latent distribution Ta = EodPaat Introducing weightings Xe, dd 2 0, determining which mapping we regularize, this remains to minimize
min Ex Ex-pum xx (x/x) log pe (xiz) + 8DxL (gx (z)
+ GMG,, (es) + AaGM (de),
0.6
① reconstruction
(ii) prior matching
(ii)) geom. preservation
66
7

---

Published as a conference paper at ICLR 2025
Source p
Target
No reg.
+DST (T)
GT)
Source p
Target q
No reg.
+DST (T)
+GMG? (T)
Figure 2: Learning of Gaussian source distribution p and a uniform or circular target distribution we minimize a fitting loss together with gcometry-preserviee setup and cost. As costs on the top line, we use L2 = || points while on the bottom line, we use Cos cos-sim(.,,) +), to preserve angles.
maps with the DST and the GMG. Provided mixture of
regularization. Each line correspond a different data
2 to preserve distances between
With this loss, prior matching and geometric desiderata can efficiently combined, as terms (ii) and (iii) can simultaneously be 0. Note that this loss can be easily extended more advanced prior matching objectives, such as the 3-TCVAE loss (Chen et al. 2018)) and can be combined with other regularizers, such as the HFS (Roth et 2023). We explore this strategy in experiments 4.2
4 EXPERIMENTS
4.1 GMG vs. DST: ILLUSTRATIVE EXAMPLE
In Fig. 2. we start by illustrating the difference between the DST and GMG, We train Ta to map from 3D source p to a 2D target q (first and second column) by minimizing L()) S. (Todp,g) q) R(T,), where S. is the Sinkhorn divergence (Feydy et al. 2019). We compare three settings, [i] no regularization R [ii] transported distribution To Totp after training. Without regularization (third column), we fit the marginal constraint Toiy g but not preserve the geometric features. With the DST (fourth column), we preserve geometric features but do not fit the marginal constraint, i.e., Totp # On the other hand, with the GMG (fifth column), we get the best compromise by approximating Gromov-Monge map: we fit Teip- g. while preserving the geometric features as fully possible.
DST, as regularizer and [mi]
GMG,. For all cases, we plot the
4.2 LEVERAGING THE GROMOV-MONGE GAP FOR DISENTANGLEMENT
Experimental Setup. Having demonstrated how the GMG enables (i) fitting a marginal constraint while (ii) preserving geometric features much as possible, we now apply it to disentangled representation learning. Our primary goal is to investigate whether the GMG results in enhanced disentanglement compared to the DST by efficiently combining (i) prior matching with (ii) geometric constraints on the representation space. Moreover, we aim to determine which cx, Gy to choose and what part of the pipeline should be regularized, the encoder cg. or the decoder de.
Baselines. We use the standard 3-VAE and 3-TCVAE as our starting models, with the option to apply the recent HFS regularization (Roth et al. 2023) to each, resulting in total of four base configurations. Note that the latter does not leverage geometric constraints; it is only used to enhance prior matching. We then investigate the effect of various geometry--peserving regularizations disentanglement on top of these four base configurations. We consider GMG, DST and the Jacobian- based (Ja e) regularization (Lee al. 2022) discussed in 2.1 Given the inclusion of the DST, naturally consider Nakagawa et al. (0023 as a baseline.
Metrics. We evaluate the learned representations using the DCI-D (Eastwood and Williams 2018 as it was found that it is the metric most suitable for measuring the disentanglement (Locatello et al. 2020 Dittadi al. 202T). We report mean and standard deviation over 5 seeds.
8

---

Published as a conference paper at ICLR 2025
Table I: Effect of different regularization on disentanglement (DCI-D on Shapes3D). We highlight the best method per regularization type (L22² ScL22, or Cos), and the best/second best per column.
8-VAE
B-TCVAE
B-VAE HES
3-TCVAE HFS
Base
65.8
15.6
75.0 +34
88.1
+7.4
90.2 +7.5
L2
cx ey
=
-1.---1 112
Enc-DST
59.6 +.9
75.7 +3300
88.7 #7.1
90.3 +7.9
+ Enc-GMG
62.3 =844
75.4 +5.3
88.4 +7.7
90.1 +4.33
+ Dec-DST
71.5 13.6
75.8 16.6
92.1 19.7
90.9 17.6
+ Dec-GMG
72.0 18.5
78.9 #5.0
92.5 #4.4
91.7 16.0
ScL22 cx
||
12 cy all
11 02 >
0 learnable
+ Jac
61.4 1:1.8
76.7 1:4.5
90.5 43.8
91.5 15.6
+ Enc-DST
65.8 1:1.9
73.0 +7.9
92.4 +377
89.2 +3.8
+ Enc-GMG
65.1 =5.5
76.1 #7.7
90.8 =9.
92.0 #5.3
+ Dec-DST
67.4
77.9 14.5
93.2 19.7
94.5 #6.9
Dec-GMG
70.0 55.9
81.0 #3.2
93.3 +8.6
96.1 #3.8
Cos CX = сy
cos-sim(;)) (:+)
+ Enc-DST
69.2
#9.1
77.2 +7.5
87.7 #7.7
90.5 55.9
+ Enc-GMG
70.9 99.5
79.6
16.6
92.5 +5.9
93.5 16.9
+ Dec-DST
76.8 4.11
81.3 14.7
87.5 #3.3
91.9 19.4
+ Dec-GMG
82.1 4.5
83.7
+8.8
95.7 #5.8
96.9 4.9
Datasets. We benchmark over four 64 x 64 image datasets: Shapes3D (Kim and Mnih, 2018).
DSprites (Higgins et al. 2017). SmallNORB (LeCun al. 2004), and Cars3D (Reed et al. 2015).
Hyperparameters. To ensure a fair experimental comparison, we follow recent works (Locatello et 2019a| 2020) Rooth et al. 2023) by using the same architecture and hyperparameters. We perform a similar small grid search for the weighting terms: 3 for the KL loss and y for HFS, Additionally, we include the weighting for the geometry-preservine regularizer in the grid search. Note that we search over the same loss weightings À for DST, GMG, and Jac. For all experiments with the GMG, we compute it with Alg. []] and systematically use Eg = 0.1 and stat mean. As result, is not included in the grid search. We conduct ablation study on en in App. D.3 We use batch size = 64. this scale, the computational cost of compute the GMG loss for a batch negligible, about 3 milliseconds. See App. for full details on hyperparameters.
Which costs cx, Cy should we choose? The first question that naturally arises when using a peosetry--reserving regularizer is: Which geometric features should be preserved? Previous works (Lee et al., 2022, Nakagawa et al., 2023) Huh et al. 2023) focused on preserving scaled distances between points, with the scale being learnable. We follow and extend this approach by also investigating plain distances and angles. This leads to three choices for cx, Cy (L22 ScL22, Cos), introduced in 2..] following the hierarchy of eeometry-preserviee (2022). We benchmark these on Shapes3D across various settings including DST, GMG, and Jac. Table shows that angle preservation (Cos), previously unconsidered for disentangled representation learning, consistently outperforms (scaled) distance preservation. This result is intuitive, as preserving angles imposes a weaker constraint, allowing for greater latent space expressiveness. In practice, preserving scaled distances seems to overly restrict the expressiveness of the latent space.
mappings proposed in Lee et al.
Should we regularize the encoder Eg., or the decoder dg? The next question we aim to answer is whether the decoder or encoder should be regularized. Therefore, we follow the previous setup on Shapes3ID and benchmark all geometry-preserving regularizers on eg and dg as reported Table We find that regularizing the decoder is beneficial over regularizing the encoder. We hypothesize this is due to the regularization of dg offering a stronger signal as its gradients impact both Φ and as in this case, the reference
Cofipaan is the distribution of encoded images. Our findings align with prior works (Lee et al.) [202]] Nakagawa et al. 2023), which focus on regularizing de yet do not offer this type of analysis. Additionally, we find that the GMG consistently outperforms the DST over all costs. Overall, the GMG on de with Cos achieves best DCI-D results over all baselines. Consequently, moving forward we regularize the decoder for angle preservation.

---

Published as a conference paper at ICLR 2025
Table 2: Effect of GMG and DST with Cos as costs on disentanglement, as measured by DCI-D, over three datasets. We highlight the best, and second best result for each dataset and method.
With Cos costs
B-VAE
B-TCVAE
B-VAE HFS
B-TCVAE + HFS
DSprites (Higgins et al.
2017
Base
26.2 18..
32.3 19..3
33.6
17.9
48.7 100.2
+ Dec-DST
28.6 19.3
32.4
88.5
39.3
=1811
49.0 111.2
+ Dec--GGG
39.5 15.2
42.2 3.6
46.7 2.0
50.1 885
SmalINORB (eCun et al.
2004)
Base
26.8 +:022
29.8
+0.4
26.8
+:02
29.8 #0.4
+ Dec-DST
28.2 +0.3
29.9
00.4
28.2
+033
29.9 00.4
+ Dec-GMG
28.3 +0.6
29.9 0.5
28.3 10.6
29.9 0.5
Cars3D (Reed et al.
2015)
Base
29.6 +5.7
32.3 #4.6
29.6
+5.7
32.3 #4.6
+ Dec-DST
26.8 +3.6
33.7 #4.2
26.8
336
33.7
#4.2
+ Dec-GMG
30.1 +5.6
36.4 =5.7
30.1 55.6
36.4 =5.7
GMG Consistently Enhances Disentanglement. To further validate our findings, we benchmark the GMG for decoder regularization with angle preservation (Cos) against its distortion counterpart across three more datasets. We report full results in Table 2 Again we observe that the GMG outperforms performs equally well to its distortion equivalent. Note that for SmalINORB and Cars3D, we found no benefits with respect to DCI-D in adding HFS and obtained the best results without it. We emphasize that using the GMG with Cos regularization significantly improves results for all datasets. This establishes the GMG as an effective tool for enhanced disentanglement.
Stability of the GMG. Furthermore, we investigate the stability of the GMG com- pared to the DST. Computing the GMG using Algorithm [ involves solving an op- timization problem through GW solver.
1.0
Cos-Sim between Gradients
0.8
This experiment aims to demonstrate that our proposed method for solving the GW
0.6
Regularizer
problem leads to a stable loss in the
DST
GMG We assess the stability by mea- suring the alignment of GMG gradients V,GMG,, (Ta) across 5 randomly sam- pled batches from each dataset In ~ D, for a fixed neural map Ta. We repeat this procedure for each of the four datasets D considered and apply the same methodol- ogy the DST. Figure presents results. We observe that the GMG's gradients ex- hibit significantly higher alignment com- pared to those of the DST, demonstrating greater stability of our proposed regularizer. This suggests that the GMG, which accounts for the the minimal distortion, effectively mitigates the inherent variability inherent of the distortion, leading to more stable gradient computations.
0.4
GMG
Shapes3D DSprites
Cars3D
SmallNorb
Figure 3: Stability analysis for the DST and the GMG applied to the decoder de with cost Cos. For each regu- larizer, we assess stability by measuring the alignment of its gradients across 5 batches. We compute the cosine similarity between all pairs of gradients. Values closer to 1 indicates better alignment, hence better stability.
5 CONCLUSION
In this work, we introduce an OT perspective on unsupervised disentangled representation learning to incorporate general latent geometrical constraints. We derive the GMG, a provably weakly convex regularizer that measures whether map T transports a fixed reference distribution with minimal distortion of some predefined geometric features. By formulating disentangled representation learning as transport problem, integrate the GMG into standard training objectives, allowing for incorporating and studying various geometric constraints on the learned representation spaces. We show significant performance benefits of our approach four standard disentanglement benchmarks.
10

---

Published as a conference paper at ICLR 2025
6 REPRODUCIBLITY
In this work, we introduce the GMG, computed as detailed in Algorithm ! To facilitate reproducibil- ity, we provide the implementation code for computing the GMG on source and target batches in Appendix E Comprehensive proofs, including all underlying assumptions, are presented in Ap- pendix B For our experiments on disentanglement benchmarks, we adhere to standard practices, employing streamlined preprocessing across all datasets. Detailed descriptions of these procedures are available in Appendix C All experiments described in this paper can be conducted using a single RTX 2080TI GPU, ensuring accessibility and replicability of our results.
ACKNOWLEDGEMENTS
Co-funded by the European Union (ERC, DeepCell 101054957). Views and opinions expressed are, however, those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council. Neither the European Union nor the granting authority can be held responsible for them. Luca Eyring and Karsten Roth thank the European Laboratory for Learning and Intelligent Systems (ELLIS) PhD program for support. Karsten Roth also thanks the International Max Planck Research School for Intelligent Systems (MMPR--IS) for support. Zeynep Akata was supported by BMBF FKZ: 011S 18039A, by the ERC (853489 DEXIM), by EXC number 2064/1 project number 390727645. Fabian J. Theis consults for Immunai Inc., Singularity Bio B.V., CytoReason Ltd, Cellarity, and has ownership interest in Dermagnostix GmbH and Cellarity
REFERENCES
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives, 2014. [
Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexander Lerehner. Towards a definition of disentangled representations, 2018. T
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of dis- entangled representations. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Pro- ceedings of the 36th International Conference on Machine Learning, volume 97 of Proceed- ings of Machine Learning Research, pages 4124 PMLR, 09-15 Jun 2019a. URL 26
ht https://proceedingg wwweeeiiinnsm.....oommm www.pp.ess///e..m..e htt/1//ate11l1111////// 9a. html
Francesco Locatello, Ben Poole, Gunnar Raetsch, Bernhard Schölkopf. Olivier Bachem, and Michael Tschannen. Weakly-supervise disentanglement without compromises. In Hal Daum III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 6348 6359. PMLR, 13-18 Jul 2020. URL
ml wtt.sp/els/slcceeee... http://ccc..aana.pprrrrse////aaa
html
26
Karsten Roth, Mark Ibrahim, Zeynep Akata, Pascal Vincent, and Diane Bouchacourt. Disentan- glement of correlated factors via hausdorff factorized support. In The Eleventh International Conference Learning Representations, 2023,
id=OkcJhpQiGix
26
Kyle Hsu, Will Dorrell, James C. R. Whittington, Jiajun Wu, and Chelsea Finn. Disentanglement via latent quantization. Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Vitria Barin-Pacela, Kartik Ahuja, Simon Lacoste-Julien, and Pascal Vincent. On the identifiability of quantized factors, 2024. I
Locatello, G. Abbati, T. Rainforth, S. Bauer, B. Schölkopf, and O. Bachem. On the fairness of disen- tangled representations, In Advances in Neural Information Processing Systems 32 (NeurIPS 2019), pages 14584-14597, Curran Associates, Inc., December 2019b. URL rness
nips httpppaper/960000000 hrt.://eeeee..t..ooeooeoooaeaaee htt.r//nnnnnatoooeoeeeeeeeaaaaoa

---

Published as a conference paper at ICLR 2025
Frederik Träuble, Elliot Creager, Niki Kilbertus, Francesco Locatello, Andrea Dittadi, Anirudh Goyal, Bernhard Schölkopf. and Stefan Bauer. On disentangled representations learned from correlated data. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10401-10412. PMIR, 18 24 Jul 2021. htRlhttps: pettps://proceeet ngs. ml press/ v1 v139/traublel1e. html
Michal Rolinek, Dominik Zietlow, and Georg Martius. Variational autoencoders pursue pea directions (by accident). In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2019. 1
Dominik Zietlow, Michal Rolinek, and Georg Martius. Demystifying inductive biases for (beta- >aee based architectures. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 12945-12954 PMLR, 18 24 Jul 2021. /ttpttts:://rrrcediiii mlr
Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
2
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations, 2017.
10 27
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 2649-2668. PMLR, 10-15 Jul 2018. URL
https::/prroceet ngs ml press. v i ml 9b. html
Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. VARIATIONAL INFERENCE OF DISENTANGLED LATENT CONCEPTS FROM UNLABELED OBSERVATIONS, In International Conference on Learning Representations, 2018. URL [RL[iattpp://peereviie fott:/foommmii=illle
Christopher P. Burgess, Irina Higgins, Arka Pal, Loc Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in beta-vae. CoRR, abs/1 804.03599, 2018.
URL httt:://aaxxiii org/ abs//804 03599
Ricky T. Q. Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. Isolating sources of disentanglement in variational autoencoders. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL htt.:h:ss////oceeeeseeeeeeeeeeee ngs neur ps CC/ pdf
Amos Gropp, Matan Atzmon, and Yaron Lipman. Isometric autoencoders, 2020.
Nutan Chen, Alexej Klushyn, Francesco Ferroni, Justin Bayer, and Patrick van der Smagt. Learning flat latent manifolds with vaes, 2020a.
Yonghyeon Lee, Sangwoong Yoon, MinJun Son, and Frank C. Park. Regularized autoencoders for isometric representation learning. In International Conference on Learning Representations, 2022.
URL httltttps///pperee//// net/ forum? 1ddmmxxxx171100
Daniella Horan, Eitan Richardson, and Yair Weiss. When is unsupervised disentanglement possible? In A. Beygelzimer, Y. Dauphin, P. Liang, Jnd J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL
XeEF 9riB93S
3
Nao Nakagawa, Ren Togo, Takahiro Ogawa, and Miki Haseyama. Grosov-wasserstein autoencoders, 2023.
12

---

Published as a conference paper at ICLR 2025
In Huh, changwook jeong, Jae Myung Choe, Young- Houg-GGu-e Kim, and Dae Sin Kim. Isometric quotient variational auto encoderss for structure- preserving representation learning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023, URL http:://penneeiier
net//formm d= d=3dggb3ggR4
9
Jaehoon Hahm, Junho Lee, Sunghyun Kim, and Joonseok Lee. Isometric representation learning for disentangled latent space of diffusion models,
11451
Filippo Santambrogio. Optimal Transport for Applied Mathematicians. Birkhäuser, NY, 55(58-63): 94, 2015. 2 19
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport. Foundations and Trends in Machine Learning, (1(6-6), 2019, ISSN 1935-8245,
Karl- Tarl-Theodor Sturm. The space of spaces: curvature bounds and gradient flows on the space of metric measure spaces, volume 290. American Mathematical Society, 2023. 23
Facundo Mémoli. Gromov-wasserstein distances and the metric approach to object matching.
Foundations of computational mathematics, 11:417-487, 2011.
Tho Uscidda and Marco Cuturi. The monge gap: regularizer to learn all transport maps, 2023.
Michael Moor, Max Horn, Bastian Rieck, and Karsten Borgwardt. Topological autoencoders, 2021. URL httpt/tsss///aiii..toe http:/ab//1gg 00722
Nikita Balabin, Daria Voronkova, llya Trofimov, Evgeny Burnaev, and Serguei Barannikov. Dis- 12696
entanglement learning via topology, 2024. URL
Tjalling C Koopmans and Martin Beckmann. Assignment problems and the location of economic
activities. Econometrica: journal of the Econometric Society, pages 53-76, 1957.
Athina Sotiropoulou and David Alvarez-Melis. Strongly isomorphic neural optimal transport across
incomparable spaces, 2024. URL httphttss/xxxi..ottt http://bbs//...... 2407.14999 14957
3
Gaspard Monge. Mémoire sur théorie des déblais et des remblais. Histoire de I' l'Académie Royale des Sciences, pages 666-704, 1781. 17
Facundo Mmoli and Tom Needham. Comparison results for gromov-wasserstein and gromov-monge distances, 2022.
Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. Advances in Neural Information Processing Systems (NeurIPS), volume 26, 2013. 28
Gabriel Peyré, Marco Cuturi, and Justin Solomon. Gromov-wasserstein averaging of kernel and distance matrices. In International Conference on Machine Learning, pages 2664-2672, 2016.
28
Meyer Scetbon, Gabriel Peyré, and Mareo Cuturi. Linear-time gromov wasserstein distances using low rank couplings and costs. In International Conference on Machine Learning, pages 19347- 19365. PMLR, 2022. 28
Théo Dumont, Théo Lacombe, and François-Xavier Vialard. On the existence of monge maps for the gromov-wasserstein problem. 2022. 5
Marco Cuturi, Laetitia Meng-Papaxanthos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein. arXiv Preprins arXiv:2201. 123244 2022. 6
Damek Davis, Dmitriy Drusvyatskiy, Kellie J. MacPhee, and Courtney Paquette. Subgradient methods
for sharp weakly convex functions, 2018.
13

---

Published as a conference paper at ICLR 2025
Jonathan Weed and Francis Bach. Sharp asymptotic and finite-sample rates of convergence of empirical measures in wasserstein distance, 2017. URL htt.ttts//axx...org////oa//
00087 7
Zhengxin Zhang, Ziv Goldfeld, Youssef Mroueh, and Bharath K. Sriperumbudur. Gromov-wasserstel
distances: Entropic regularization, duality, and sample complexity, 2023. 7 28
Jean Feydy, Thibault Séjourné, François-Xavy Saavirr Vialard, Shun-Ichi Amari, Alain Trouvé and Gabriel
Peyré. Interpolating between Optimal Transport and MMD using Sinkhorn Divergences. In International Conference on Artificial Intelligence and Statistics (AISTATS), volume 22, 2019,
Cian Eastwood and Christopher K. I. Williams. A framework for the quantitative evaluation of disentangled representations. In International Conference on Learning Representations, 2018.
URL
wtwp:/eeeeeeeeeeeeeeeeee net Eorum SeerrummidByy-ooo 7d g
Andrea Dittadi, Frederik Träuble, Francesco Locatello, Manuel Wuthrich, Vaibhav Agrawal, Ole Winther, Stefan Bauer, and Bernhard Schölkopf. On the transfer of disentangled representations in realistic settings. In International Conference Learning Representations, 2021. URL
https::/openzeviiet
.et/ -
Yann LeCun, Fu Jie Huang, and Léon Bottou. Learning methods for generic object recognition with invariance to pose and lighting. Proceedings of the IEEE Computer Society Conference Computer Vision and Pattern Recognition, 2:1197-11104, 2004. ISSN 1063-6919. Proceedings of the 2004 IEEE Computer Society Conference Computer Vision and Pattern Recognition, CVPR 2004 Conference date: 27-06-2004 Through 02-07-2004,
Scott E Reed, Yi Zhang, Yuting Zhang, and Honglak Lee. Deep visual analogy- making. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, ed- itors, Advances in Neural Information Processing Systems, volume 28. Curran Asso- ciates, Inc., 2015. URL httpttt////prooeeeinnss neurips www.eeis..c.papeeeeeeeeeeee pct.c//aper/221///111ee pdf 10
I. Kantorovich. On the transfer of masses (in russian). In Doklady Akademii Nauk, volume 37, page 227, 1942. 17
Thibault Sjourn, Conoos-Xaviee Vialard, and Gabriel Peyr.
The unbalanced gromov wasserstein
distance: Conic formulation and relaxation, 2023. 19 20
Garrett Birkhoff. Tres observaciones sobre el algebra lineal. Universidad Nacional de Tucumán Revista Series A, 5:147-151, 1946. 20
D. Bertsimas and J.N. Tsitsiklis. Introduction to linear optimization. Athena Scientific, 1997. 20
Jean-Franois Le Gall. Inigration, Probabilits el Processus Alatoires.
21
Gabriel Rioux, Ziv Goldfeld, and Kengo Kato. Entropic gromov-wasserseiin distances: Stability,
algorithms, and distributional limits, 2023. 21 28
Tudor Manole and Jonathan Niles-Weed. Sharp convergence rates for empirical optimal transport with smooth costs. The Annals of Applied Probability, 34() February 2024. ISSN 1050-5164. doi: 10.1214/23-aap 1986. BRRL:ttp://dee doi, org
1986
22
Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, March 2004. ISBN 0521 833787. http:/ttt////tttt http://wwwwaa....ommmooo zon zon. com/ caan//xee//oo cacec//obioo redi rect.
ike orrrrcteee..lteeeo0o.oaaaa311332
K. B. Petersen and M. S. Pedersen. The matrix cookbook, October 2008 URL //ww2
imm. wtw..ttudddppbbbpppppppppp phpp33/ Version 200811101 23 25
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, lurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, John Quan, George Papamakarios, Roman Ring, Francisco
14

---

Published as a conference paper at ICLR 2025
Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan,
Luyu Wang, Wojciech Stokowiec, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL
http://githu>: com/ deepmind
25
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations (ICLR), 2014. 26
Jason Altschuler, Jonathan Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration, 2018. URL
09634
28
Zoe Piran, Michal Klein, James Thortoon, and Marco Cuturi. Contrasting multiple representa- tions with the multi-marginal matching gap, 2024. [tRL:tttss///axx...crr/////////
19532
28
Aram-Alexandre Pooladian and Jonathan Niles- Weed. Entropic estimation of optimal transport maps.
arXiv preprint arXiv:2109. 2021. 29
Parnian Kassraie, Aram-Alexandre Pooladian, Michal Klein, James Thornton, Jonathan Niles-
Weed, and Marco Cuturi. Progressive entropic optimal transport solvers, 2024. URL https
//arxiv.org/ann abcgg//bs/2// 05061
29
Dominik Klein, Tho Uscidda, Fabian Theis, and Marco Cuturi. Entropic (gromov) wasserstein flow
matching with genot, 2024. 29
Chris Burgess and Hyunjik Kim. 3d shapes dataset. http:///gthhbbooomooooom//////ss 2018. 3T
Andrea Burns, Aaron Samna, Dilip Krishnan, and Aaron Maschinot. Unsupervised disentanglement without autoencoding: Pitfalls and future directions. arXiv preprint Stti::2000.066613, 2021. 31
Julius von Kgelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schlkopf, Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. In Advances in Neural Information Processing Systems, 2021. 31
Cian Eastwood, Julius von Kügelgen, Linus Ericsson, Diane Bouchacourt, Pascal Vincent, Mark Ibrahim, and Bernhard Schölkopf. Self-supervised disentanglement by leveraging structure in data augmentations. In Causal Representation Learning Workshop at NeurIPS 2023, 2023. URL
htttps://openrviile wtt.leewwwe/ooorriiiiiiooo (osoorummmdddd11sspo@8111111 31
net
Stefan Matthes, Zhiwei Han, and Hao Shen. Towards unified framework of contrastive learning for disentangled representations. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL
31
Laurence Aitchison and Stoil Krasimirov Ganev. InfoNCE is variational inference in recognition parameterissed model. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL
hnttps://openrevilt
f wtt.frrrmmmmidccchhwowwwawaaaaa 31
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Hal Daum III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Pro- ceedings of Machine Learning Reseurch, pages 1597-1607. PMLR. 13-18 Jul 2020b. URL mlr. www.epresssse.e. /alass/vvlie
html
31
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stephane Deny. Barlow twins: Self- supervised learning via redundancy reduction. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning. volume 139 of Pro- ceedings of Machine Learning Research, pages 12310-12320. PMLR, 18-24 Jul 2021. URL
1 www.press/v.. 39/ 2tt99//ootar21ea. htm
31
Adrien Bardes, Jean Ponce, and Yann LeCun. VICReg:
regularization for self-supervised learning. In International Conference on Learning Representations, 2022. URL httpps://openrse net http///oommmmiimmmmmmmo202
31
15

---

Published as a conference paper at ICLR 2025
Quentin Garrido, Yubei Chen, Adrien Bardes, Laurent Najman, and Yann LeCun. On the duality between contrastive and non-contrastive self-supervised learning. In The Eleventh International Conference Learning Representations, 2023. URL https: / appenreview.net/forree
Id=kDEL.91Du
31
16

---

Published as a conference paper at ICLR 2025
APPENDIX
The Appendix is organized as follows:
Section
provides additional background information to supplement the main text.
Section
presents all theoretical proofs, including detailed assumptions.
Section
outlines comprehensive experimental details, ensuring reproducibility.
Section
offers supplementary empirical results that further support our findings.
Section
includes the implementation code for computing the GMG.
A ADDITIONAL BACKGROUND
A.1 REMINDERS ON MONGE AND KANTOROVICH OT
In this section, we recall the Monge and Kantorovich formulations of OT, which we will use to prove various results. These the classical formulations of OT. Although introduce them here after discussing the Gromov-Monge and W formulations, should be noted that they are generally introduced beforehand. Indeed, the Gromov-Mong and Gromov- Wasserstein formulations were historically developed to derive OT formulations for comparing measures supported on incomparable spaces.
Gromov-Wasserstenn
Monge Formulation. Instead of intra-domain cost functions, we consider here an inter-domain continuous cost function e: X' y This assumes that we have a meaningful way to compare elements x, from the source and target domains. The Monge (1781) problem (MP) between p E ( and p P()) consists of finding map I' ; X + that push-forwards p onto p, while minimizing the average displacement cost quantified by
inf / c(x, dp(x)) T:Thp=
(MP)
We call any solution T" to this problem a Monge map between and for cost c. Similarly to the Gromov-Monge Problem Problem[GMM) solving the Monge Problem is difficult, as the constraint set not convex and might be empty, especially when p, are discrete.
Kantorovich Formulation. Instead of transport maps, the Kantorovich problem couplings E (Tp,
KPP) seeks
W(p. q) == min / c(x,y) (x..
(KP)
#Ell(p.4) xxxy
An optimal coupling x¹ solution of (KP) always exists. Studying the equivalence between (MP) and (KP) is easier than in the Gromov-Monge and feasible, the Monge and Kantorovich formulations coincide and क = (Id,T")tp.
cases. Indeed, when (MP)
is
A.2 CONDITIONALLY POSITIVE KERNELS
In this section, recall the definition of a conditionally positive kernel, which is involved in multiple proof' relying on the linearization the Gromov- Wasserstein problem as a Kantorovich problem.
Definition A.I. A kernel A: Rdd x Id + IR is CPD, i.e., conditionally positive, if it is symmetric and for any X1, Xn Rd and Rn" s.t. 1n O, one has
a,a, k(x;, x,) 2 0
CPD include all positive kernels, such as the inner-product kixx, x') (x,x',, or the cosine similarity k(x,x') x') = cos-sim(x,, x') = | xx but also the (scaled) negative squared Fuclidean distance x(x,x') -a x1112, positive kernel - for the inner product and the cosine distance or its opposite - for the squared Fuclidean distance. Additionally, CPD kernels also include more exotic cost functions, such as:
O. Therefore, each of the costs of interest is either conditionally
17

---

Published as a conference paper at ICLR 2025
The RBF kernel x(x,x') = -
The power kernel x(x,x') =
00p00pp>2
The thin plate spline kernel kxx, x') 1xxx'111oogl X-)
. The inverse multi- qultii-uuadraic kernel k(x,x') = -
As a result, the family of CPD kernels includes a large variety of cost functions, which can be used to define various GMGs.
B PROOFS
B.I THE GMG CHARACTERIZES GROMOV-MONGE OPTIMALITY
We show here that if GMG, (T) O, then T is a Gromov-Monge map between and Tür for costs cx, Cy. As the set of deterministic couplings {TT == (Ia, Fr Ftr = Ttr C [1 Tr we immediately get that
[ (cx(xxx))cxx((xx)xx..)))))))))> cx((((), drx) GW(r,
inf Fdr- Thr xxxX
(7)
On the other hand, if GMG,(T) = 0, one has
GW Ttr) xxxX
- -07x
(8)
Therefore, combining Eq. 7 and Eq.
we get
inf Fir=Thr Jxxx I (cx(x, x') (F((F))) drx))))
(9)
-
xxxx (xx(xx)))<((((())) (xx) dr(x) dr(x')
Finally, as ' naturally satisfies the marginal constraint, we conclude that 7' is a Gromov-Monge map between and Thr for costs cx, Cy-
B.2 ON RESCALING THE COSTS MATRICES IN THE ENTROPIC GW SOLVER
We remind, from Eq. (EGWP) that
GWe 9n)
min PEU,, 7.
Σ (Cx,,
CC, "P. PP.P
(10)
By developping each terms, and using the fact that Uns we get
GWe(Pns = min 1 (C 1n 1n) 1 (CS
1 1,) 2(CxPCy, - =P)
PEU,
(11)
1 (C In, 1,) + n2 (C) 1n 1,)
min PEU,
-C P) - cH
where C Cx O Cx. with @ the Hadamard (i.e., elementwise) product, and similarly for C As can see that the two terms on the left do not depend on P, they do impact the minimization, an OT coupling p solving the problem satisfies:
P E arg min -2 PH)e
(12)
PeU,
As result, if we now replace Cx and Cy by Cx/stat(Cnn (Cx) and Cy/stat(Cy)) respectively, the new OT coupling p solving the problem satisfies
Cx P Cy
D E arg min
-2 stat(Cx)) stat(Cyy) PHH
PEU,
2
+> P arg min PEU,
stat(Cxx) stat(Cy)) (Cx PC, P) s(((P)
(13)
D E arg min -2CCxPC,, P) - stat(Cx)) stat(Cy)) {HH(P)
PEU.
18

---

Published as a conference paper at ICLR 2025
where in the last line, the fact that stat stat(Cx)) stat(Cy) 0. This yields the desired equivalence on scaling the cost matrix and adapting the entropic regularization strength.
B.3 POSITIVITY OF THE ENTROPIC GMG ESTIMATOR
Recall that
GMG,,,e DSTr, (T) GWe(r Ttrn)
DSTr, (T)
min PEU
Σ
(cx(x;,x;)- (o(x)))))) - eP)
i.",",''==1
For any coupling Un since (PH(P) - Eij=1
Pi Poglog(P1)) ) 0, one
(ox(x) f
Σ
Poy
As result, applying minimization both sides yields that GW Tirm) GWo(ra, Tirn) GW Tfrn). and therefore:
GMG,....T > GMG,,..o(T) GMG,, (T) 2 0.
B.4 PROOFS OF PROP. 3.2
Proposition 3.2. If GGGG, = O, Vs E P(X) s.t. Spt(s) C Spt(r), one has GMG,(T) = 0.
== (Id, T) is an optimal Gromov-Wasee Wasserstein coupling, solution of Problem GWP))etween and Tir for costs and cy. Therefore, from (Sjourn et al. 2023 Theorem. 3), n an optimal Kantorvich coupling, solution of Problem KP) between and Tür for the lincarized cost:
Proof. Let r,, & as described and suppose that GM = 0. Then, ==
1 ē: (x, (x EX [ 2 _xx(xx
d (x',''')))
(14)
Additionally, x y' is a compact set as a product of compact sets, so since (x, y) + lxxxxxxxx101011 x') yy(y,y') continuous cx and cy are continuous, it is bounded on X' x V. Afterward, since x² has finite mass, by Lebesgue's ' dominated convergence Theorem, it follows that č is continuous, and hence uniformly continuous, again since X' x y is compact.
Afterwards, by virtue of (Santambrogio 2015 Theorem 1.38), Spt (T*) ) is a č-cyclically monotone (CM) set (see (Santambrogio 2015 Definition. 1.36)). From the definition of cyclical monotonicity, this property translates to subsets. Then, by defining л- = (Id. T)ds, as Spt(p) C Spt(r), one has Spt(m*) = (Id, T)ss) C Spt((Id, ())) Spt(x"). so Spt(x*) is ?-CM. Finally, since X' and y are compact, and ē is uniformly continuous, the č-cyclical monotonicity of its support implies that the coupling 7 is a Kantorovich optimal coupling between its marginals for cost thanks to (Santambrogio 2015) Theorem 1 1.49). By re-applying (Sjourn et al. 2023 Theorem. 3), we get that x vol the other words, x² = (Id, Tsss is Wasseertenn optimal coupling between s and Tps so T is Gromov-Mone map between 8 and Tps and GMG,(T)
problem between its marginals for costs cx and cy. In
B.5 PROOFS OF THM. 3.4
Theorem 3.4. Both GMG?, and "MG.")) well their finite sumple versions, are weakly convex.
Finite sample. We note X E Rund the matrix that stores the Xú, i.e. the support of Yno as rows. Then, (i) GMG?, and (ii) GMG are respectively (i) 72.a and (ii) "Jimer,, -weakly conver, where: "/mmer. >maxx (_XXT) Xmin( ( and 7/2, Yomeer, + max.=1.... ||x
Asymptotic. GMG? and (ti) GMG are respectively (i) 7/2 and (ii) Weakly convex, where: "iunerr \max (Ex-r [xx and 12 = Timner + maxxeSpt(r) ||x|
Before proving Thm.
3.4 we first demonstrate technical results that will be useful later.
19

---

Published as a conference paper at ICLR 2025
Reformulation of the empirical GMG using permutations. We start by showing that GMG,,, (T) always the sub-optimality gap of in Prob. [GMP)betweenn between in and Ttr, This occurs because Prob. (GMP) and Prob. (GWP) coincide when applied between empirical measures on the same number of points. In other words, we can reformulate Prob. (GWP) between Ia and Tprn using permutation matrices, instead of (plain) couplings.
Proposition B.1. The empirical GMG reads
GMG,, (T) = DST (T)
min Σ (cx(xx 5,=1
- (cyT(xo))) ((xo))))
(15)
oES,
Proof. We first show kernels (see A. the discrete GW couplings between uniform, empirical distributions supported on the same number of points, permutation matrices.
general results, stating that when ex, cy conditionally positive
Proposition B.2 (Equivalence between Gromov-Monge and Gromov- Gromov-Wassersee Wasserstein problems in the discrete case.). Let Pv. - - 8x, and qu - - dy. two uniform, empirical measures, supported the same number of points. We denote by Pn (P E RaxR 3o Suy Pig == ( the set set of permutation matrices. Assume that cx and ey (or -cx and -cy) are conditionally positive kernels (see[A.]). Then, the GM and GW formulations coincide, in the sense that we can restrict the GW problem to permutations, namely
GW Pn) PEUn
min
(
min n² PEPa in
Σ
(16)
min 22 S.S 1
Σ
(cx(xx, -cy(7(())
Proof. Let P E Un solution of the Gromov-Wassene between Pn. and Pr. i.e.
P E arg min PEU,
Σ (cx(x,, x;) Pr. 6'
that always exists by continuity of the GW objective function on the compact Un. We show that P can be chosen as (rescaled) permutation matrix without loss of generality.
As we assume that cx and Cy (or -cx and -cy) are conditionally positive kernels, from (Sjoun) al. 2023 Theorem. 3), P also solves:
X2 Σ
P E arg min QEU, (J
-xyx) Qey
(17)
We then define the linearized cost matrix Č' Raxn
(cx(x; -y -cy,,))))))()))))))))))
which allows us to reformulate Eq.
17 as
P arg argmin (C, Q)
(18)
QEU,
Birkhoff s theorem states that the extremal points of Un, are the permutation matrices P,. Moreover, seminal theorem of linear programming (Bertsimas and Tsitsiklis 1997 Theorem 2.7) states that the minimum of linear objective on a bounded polytope, if finite, is reached at an extremal point of the polyhedron. Therefore, P solves Eq. (8), it is an extremal point of U,, it can always be chosen as permutation matrix. Therefore, the equivalence between GW and GM follows.
20

---

Published as a conference paper at ICLR 2025
To conclude the proof of Prop. B.1 we simply remark that in Ei= dx: and Ttr,
Eis1 (xx),) uniform, empirical distribution, and supported on the number of points.
Consistency of the empirical GMG. We continue by proving consistency result for the empirical GMG, which we will later use to deduce the asymptotic weak convexity constant from the finite- sample case.
Proposition B.3. For both cx cy -1. (T), almost surely.
112 and cx =cy (,, one has GMGr, (T)
Proof. We first note that the empirical estimator of the distortion is consistent, as both costs con- tinuous, and X' is compact. We then need to study, in both cases, the convergence of GW(ra, Tirm) to GW (r,, Ttr).
To that end, we first remark that as, almost surely, ru - r in distribution, one also has that, almost surely, Tprn - Tpr in distribution. Indeed, since y is compact, T is bounded so for any bounded and continuous f y + R and X ~ r, T(X) is well defined and bounded so integrable. Afterwards, one can simply adapt the proof of the almost sure weak convergence of empirical measure based on the strong law of large numbers to show that, almost surely, Tirn - Ttr in distribution. See for instance (Le Gall Theorem 10.4.1).
We start with the squared Euclidean distance. As, almost surely, both I' + r and Tirn + Trr in distribution, the results follows from (Mémoli 2011 Thm 5.1, (e)).
We continue with the inner product. As noticed by Rioux et al. (2023 Lemma 2)in the first version of the paper the GW for inner product costs can be reformulated as:
f (x, x') dp(xx) dp(x') / yxx (y,y') dgly) dg(y")
(w
xxxx
(19)
/
+ min min MeM rell(p.4) XxX -4(Mx, dT(x,)) M
where we define M = = M/22 xdy with M = Sx ||x 12 dp(x) Sy lly dg(y). In par- ticular, they show this result for the entropic GW problem with O, but their proof is also valid for E 0. The above terms only involving the marginal, i.e., not involved in the minimization, are naturally stable under convergence in distribution, as X' and y are compact, so as X' x X and y x y As result, only need to study the stability of this quantity under the convergence in distribution of the following functional:
/
F(p.g) g) = min min MEM Jxxy
4|
(20)
We first remark that:
(F(p.)) - F(pn,gn) gn)l
I /xxx
I sup
min MeM rETT(pa))
44 y) dx(xxy) min TELl(p.T) Xxx
-4(Mx, y) dr(x,
/
I sup min MEM sell(pa)) /xxy M| - y113 di dx(,)) min mEll(pm.nn) Sxxy
2 - y11 di y)2
+2 sup MEM S 1Mxx/2dp6 / [IMx([dppxxxx)
+2 1/8 ||y 11 122 dgly) -/r
(21)
Then, show the convergence of each term separately.
21

---

Published as a conference paper at ICLR 2025
For the first term, remark that (up constant factor) it can be reformulated:
(M q) W 9n))
sup
MEM
where we remind that that W3 is the (squared) Wasserstein distance, solution of Eq. (KP) induced
by c(x, y) ||x yl1.- By virtue of (Manole and Niles- 2024 Theorem 2), there exists
constant C > O, s.t. we can uniformly bound
W g) (MMpn qn) S Cn 1/d
sup MEM
and the convergence follows.
For the second one, this follows from from the convergence in distribution of pn. palong with the Ascoli-Arzela theorem, since both M and X are compact sets, so the {Sm I' x + |Mx]]] are uniformly bounded and equi-continuous.
For the third one, this follows from the convergence in distribution of qo to
As result, we finally get GW (pr, Qn) (W
Weak convexity. Finally, we demonstrate some useful results weakly convex functions on Rd. Definition B.4. A function f Idd + R is y-weakly convex if f +
|| -12 is convex.
From the definition, we see that if ff is y-weakly convex, than f is also y' weakly convex for any Y' 2 y. This naturally extends weakly convex functionals F on L2(r).
Lemma B.5. Let A (R) symmetric matrix and define the quadratic form fA E Rdd x Ax. Then, JA is max(0, -Amin -A) -weakly convex.
Proof. We the fact that twice continuously differentiable function is convex i.f.f. its hessian is positive semi-definite (Boyd and Vandenberghe 2004] 1.3.1.4)). Therefore, fA is convex i.f.f. convex, i.e. O-weakly convex. Otherwise, Amia (A)l 112 has hessian A Amia (A) 2 0, so it is convex, which yields that fA is Ami -eeeelyly convex.
V IA A 2 If Amin A) > 0, then A 2 0
fA
Lemma B.6. Let (filier a family of w-eaklly convex functions, with potentially infinite 1. Then, f:xeR + supier fx(x) is 7-weakly convex.
Proof. As the fa are y-weakly convex, f, + 1y is convex, so ト supier f((x) + 1
(supier f((x)) + / is convex (Boyd and Vandenberghe 2004 Eq. (3.7)). Therefore, the y-weak convexity of follows
Lets now proceed to prove the main Thm. 3.4
Proof of Thm. 3.44. Finite sample. We first study weak convexity of GMM i.e. the Gromov- Monge gap for the inner product. For a map T L2(r), it reads
GMG (MG
(T)
x1, (T(x;), ((x)))
min PEU
xx), (T(x),
As r. and Tir, are uniform empirical supported on the same number of points, using Prop. B.2 we reformulate the RHS with permutation matrices, which yields
22

---

Published as a conference paper at ICLR 2025
GMG
1 n²
(i,
(T(x,), T(x;)))
Σ PEPn (.
xxx,,x))) ((x
From this expression, GMG can be reformulated as a matrix input function. Indeed, it only depends on the map T via its values on the support of rn, namely x Xa. X.. Therefore, we write
i, : T(x:), and define X, Raxd which contain observations x; and t, respectively, stored as rows. Then, studying GMG Pm
remains to study
n
S(T) = + Σ (x)) tJ=1
(t,, t) +
min PEñ
Σ "."""""==
(x,, - (t,, P.,
By developing each term and exploiting that for any P Pa. Pl, PT1, I we derive
S(T)
1
(x,,
(ti,
min
Σ
Xx,, xp) P.y
+J=1
i
Σ (x, tz) B Pij P Vi
Σ 11=1
(x, (tc, tj)
PEP, n²
mnax
PT
4 TT
max 1 PEP,
XX P P -X
(L(' XX XX P Xx XX
max PEP,
max PEP.
(Ax.PT, T)
where define Ax,p t P - Rax" To study the convexity of this matrix
(P
n. input function, we vectorize it. From (Petersen and Pedersen 2008 Eq. (520)), we note that, for any M E Raxa
MT. T) vec( vec(MT) vec(T) (M @ (e)v)c(T
where vec the vectorization operator, raveling a matrix along its rows, and is the Kronecker product. Applying this identity, we reformulate:
f(T) = max vec(T) ( (Ax.P & In (T) PEU,
(22)
To study the convexity of r, we study the convexity of each F'x (T) == vec(T) (A
(A ?
T, vec which are quadratic forms induced by the Ax,p & T,. This remains study the (semi-) positive definiteness of the matrices Ax,p & Ln. As each Ax,p E Ruxn is symmetric and square, Ax,p & In is also symmetric and from (Petersen and Pedersen 2008 Eq. (519)) its eigenvalues the outer products of the eigenvalues of Ax,P and Z,, namely
cig( Ax. & I,) = (A(AA.P) {T(Tn))ssssn
(Ax,)))), .... Ai (Ax.))..
(.. ...
(23)
times
follows that the minimal eigenvalue of Ax.P & In is Amin((X & 1m)) Amin (Ax.p). Utilizing the expression of Ax.P
Amin (Ax.p)
n Amin
2
n Omin P
XX T (XTP)) + Amin -XX >minn-XXX
(24)
22 Ammin P
1XX TP) Amaax (XX])
23

---

Published as a conference paper at ICLR 2025
Reminding that P Un. one has P the same eigenvalues. In particular Amin (P P) = Amin XX follows that
P-1 so P XX" and XX similar, and they have
Combining these results,
Amin Ax.P @ In) Amin (Ax,P)
2 1 nd (Auin (AminnnXXXX (XXT) Amaax (XX
(25)
We then remind that each rAx.P is the quadratic form defined by Ax,p In by applying Prop. B.5 it is Ax,P & I weeakl convex, and hence 1 (Amax DOmaa(XXX (XXT) Amin (XX J-weaklly convex. Therefore, applying Prop. (B.6). r is a (Amax XXX) - Amin(XX ---wekly convex, in Amin (XX }), r is nt Vinner weakly convex. This implies that T + f(T) + 1 Tinner | is convex. By reminding that T stores the T(x,) as rows,
R Reminding that Tinnnrr = +.2 (Amax [XX]
1 | Consequently, GMG is Tinner (r(r.)
Fn.
We then study the convexity of GMG,, We follow exactly the same approach. One has:
n
GMG, (T)
X
x,112 IT(xx; 1(3
n2
d||x
xy - - - - (Tx,))
PEP
=1
Similarly, studying the convexity of GMG?, (T) remains study the convexity of the matrix input function:
(T) =
||||
- x,113 ||t: - 1,1
1-=1
min n PEP.
Σ
2
||
x,13 - ||t, - -
As before, by developing each term, one has:
(T))
(xi, xv) (bytt))Pis
Σ
P
P
11
PEP,
2n
í.f=1
(x X²)
(ti, t2)
+
x
2
n2
2n
The quadratic terms in can be factorized before using Ax.P. For the new terms w.r.t. the inner product case, we introduce Dx diag( Xx |x11l...... |x, (12), and remark that rewrite:
P
1
| 1
vec(T) (2. (PT - I,) & Dx) vec(T)
2n
2T
1-=1
1-=1
As we can always symetrize the matrix when considering its associated quadratic form, we have:
Σ P
X 1 vec(T)) T (L (PT +P) In) 8 Dx) vec(T)
As result, we denote Bx = (P + +P) I.) 8 Dx and finally get:
(TT)
max vec(T)] (Ax.P @ In + Bx vec(T")
PEP.
24

---

Published as a conference paper at ICLR 2025
As did for f. studying the weak convexity of f remains to lower bound the minimal eigenvalue of Ax.P @ In + Bx... First, one remark that:
Amin ( @ In + Bx,P) 2 Amia Ami((A,pp (Ax, In.) + Amin (Bx.P)
As we we have already lower bounded Amin (Ax,P In) 2 1 (Amin (XX Amax XX }), we focus the RHS. Similarly, has:
Amiu (Bx.P) Amin
(P'
+ +P) 4n) Dx)
Amin
P)
C Dx) + Amin (- (-2 In @ Dx)
(26)
4x
Amiu
(in
[PT
P
& Dx
Amax (2n & Dx)
For both terms, we apply again (Petersen and Pedersen 2008 Eq. (519)). For the LHS, one has:
cig
[ + P))
( (PT +
(27)
We remark that 1 (P) +P) a symetric bi-stochastic matrix, so Amiu ( (P +P)) -. Therefore,
Amin( in (P P)) 2n As a result, since the eigenvalues of Dx are the ||x; this yields:
Amin ( (P + +P)
(xxxx
N
max |x
Similarly, we have:
Amax 2n & Dx)
max
|x
2n
from which we deduce that:
Amin (BxPP) 2 1 i= 1=1....7
max
||x:12
We then lower bound:
Amin (Ax, & In + Bx.P) 2 1 (Amin (XX - Amax (XX Amax(XX
- max
||x:1|
i= i=....,
(28)
72.
which yields the 7/22, -weak convexity of g, and finally the 72... -weak convexity of G.M...
Asymptotic. For any T. we note that, almost surely, 1 + convexity is preserved under pointwise convergence and by virtue of Prop. (B.3). we study the (almost sure) convergence of iinner. and 72.n-
|T (r) As a result, since
We start by Timeer,n- We first remark that Amax ( XXT) Amax x X). Moreover, as A E 5 (R) + Amax(A) is continuous and 1x" n X + Exxr((xx" almost surely, one has A LXXT)= + Amax (Exar [xxT] almost surely. Moreover, for any n > d, Amin ( = 0.
As a result, Tinner.n + Amax (Ex~r [xx I) almost surely, which provides the desired asymptotic result.
We continue with 72.n- We first remark that max;=1..... | I sUPxeSpt(r) ||x| As a result, by defining 72,n "inner, + maxxeSpt(r) ||x||,, GMG also 52,n convex. Moreover, maxxeSpt(r) ||x|| does not depends on n, /2.n >> Amax (Exee [xx [x) + maxxespt(r) ||x|| almost surely, which also provides the desired asymptotic result.
C EXPERIMENTAL DETAILS
All our experiments build on python and the jax-framework (Babuschkin et al. 2020). alongside ott stt-jax for optimal transport utilities.
25

---

Published as a conference paper at ICLR 2025
Table Hyperparameter grid searches for different baseline and proposed methods.
Method
Parameter
Values
3- 3-VAE
[2. 4, 8, 10, 16] [2, 4, 6, 8, 10, 16]
8 8-TCVAE
B
HFS
Y
[1, 10]
+ DST
À
[0.1, 1,5, 10, 20]
GMG
[0.1, 1, 10, 20]
Jac
[0.1, 5,
C.I. DETAILS ON DISENTANGLEMENT BENCHMARK
To effectively conduct comprehensive and representative research on disentangled representation learning, we convert the public PyTorch framework proposed Roth al. (2023) to an equivalent jax variant. We verify our implementation through replications of baseline and HFS results in Roth
al. 2023), mainting relative performance orderings and close absolute disentanglement scores (as measured using DCI-D, whose implementation directly follows from Locatello et al. (2019a) and leverages gradient boosted tree implementations from (aakit-learn)
For exact and fair comparison, we utilize standard hyperparamater choices from Roth et al. (2023) (which leverages hyerparameters directly from Locatello et al. (2019a), Locatello et al. (2020) and https: ptttss//giithi
com/ http:/oooooeeeeeeeeeae httpaarrrcdd/ieeteeeeeee anglement lib). Con- sequently, the base VAE architecture utilized across all experiment is the same as the one utilized in Roth al. (2023) and Locatello al. 2020)): With image input sizes of 64 64 Na (with N, the number of input image channels, usually 3). The latent dimensionality. not otherwise specified, is set to 10. The exact VAE model architecture is follows:
Encoder: [conv(32, 4 x 4, stride ReLU] [conv(64, 4, stride ReLUU x MLP(256), MLP(2 x 10)
MLP(256), [upconv(64, 4 4, stride 2) + ReLU] x 2, [upconv (32, 4 x 4, stride 2) + ReLUJ, [upconv(H.: 4, stride 2) + ReI.U
Decoder:
Similar, we retain all training hyperparameters from Rothh et al.
2023) and (Locatello al. 2020): Using an Adam optimizer ((Kingma and Ba 2014). 31 0.9. B2
0.999, 10 8) and learning
rate of 10-4, Following Locatello al. (2020); Roth et al. (2023) we utilize batch-size of 64, for which we also ablate all baseline methods. The total number of training steps is set to 300000.
As commonly done for this setting (Locatello et al. 2019a 2020 Roth et al. 2023), we also perform small grid search over all the hyperparameters. We report the full details in Tab. 3
For Ar and λς, we set A for the Decoder setting and Xa for the Encoder setting while altering the weighting for the other X. All experiments on single RTX 2080TI GPU.
C.2 STABILITY ANALYSIS
For the gradient stability analysis experiment, we repeat the following experiment for each of the four image datasets D that we consider. We first considered a fixed neural map Tg. which we choose be randomly initialized neural network, consisting of encoder and decoder, before any training. For &o = 1,...,5, 5, we sample a batch xf..... x ~ D and let r report the pariwise cosine similarity between the gradients of the DST and the GMG. Formally, we compute cos-sim V DSD,, (Tp), V&DST, (Ta)), and cos-sim( cos-sim(VooGee V&GMG, (Tg), VoGMG, (To)). for K, =,...., 5.
Di Set. We
D ADDITIONAL EMPIRICAL RESULTS
In this section, we report additional empirical results revolving around the regularization of encoder. First in D.I. we conduct the stability analysis when regularizing the encoder. Then, we report further
26

---

Published as a conference paper at ICLR 2025
results of encoder and decoder regularization DSprites. Lastly, take first exploratory step towards decoder-fiee free disentangled representation learning in D.6
D.1 ENCODER STABILITY ANALYSIS
We repeat the following experiment for each of the four image datasets D that we consider. We first considered fixed neural map Tg. For i = 1.....5, 5, we sample batch xi...., X,, ~ D and let r >i=1 6x; We report the pariwise cosine similarity between the gradients of the DST. Namely, for i .....,,, 5, we compute cos-sim (VDDTT (To), VoDST, (To)), and the GMG, cos-sim( VDGGM (Ta), VoGMG, r4 (To)).
1.0
Cos-Sim between Gradients
0.8
0.6
Regularizer
DST
0.4
GMG
Shapes3D
DSprites
Cars3D
SmallNorb
Figure 4: Gradient stability analysis the DST and GMG Encoder regularizations. The cosine similarity is computed between all pairs of gradients V obtained through 5 randomly sampled batches and a fixed network T9 for each dataset.
D.2 ENCODER ANALYSIS ON DSPRITES
Table 4: Disentanglement of regularizing the Encoder and the Encoder and Decoder as measured by DCI-D DSprites. We highlight best, second best, and third best results for each method and dataset.
DCI-D
3-VAE
B-TCVAE
3-VAE HES
3-TCVAE HFS
DSprites (Higgins et al.
2017
Base
27.6 13.4
36.0
15.3
38.7 +15.7
48.1 10..8
Enc-DST
32.8 15.0
36.5 +5.9
33.9 +15.9
48.9 #11.1
+ Enc-GMG
27.5
37.4 45.8
4143
45.9 110.9
Dec-DST
28.6 119.3
32.4
39.3 +18.1
49.0 111.2
+ Dec-GMG
39.5 +15.2
42.2 +3.6
46.7 +2.0
50.1 +8.5
27

---

Published as a conference paper at ICLR 2025
D.3 ABLATION ON THE ENTROPIC REGULARIZATION STRENGTH
Effect of on the entropic GW solver. In this section, we remind some insights about the entropic
GW solver (Peyré al. 2016) introduced in [.2 and used to compute the GMG in Alg. [ We provide its algorithmic details in Alg. 2 This solver naturally provides a trade-off between (i) the approximation of the true GW distance GW (p, q) and (i) the optimization speed (i.e., convergence rate). This trade-off is controlled by the entropic regularization strength E.
(i) Optimization. The solver employs a mirror descent scheme that iteratively linearizes the entropic GW problem ( and applies the Sinkhorn algorithm. Each mirror descent step corresponds to a projection with respect to the KL divergence, which can be efficiently performed using the Sinkhorn algorithm. The solver is initialized with Po zz 111 and iterates P SINKHORN ( CxP.Co Cy. e)- Since a larger e in Sinkhorn leads to faster convergence (Cuturi 2013 Altschuler et al., 2018), this subsequently accelerates each mirror descent step. In other words, it reduces the number of iterations within each Sinkhorn call, i.c., the number of inner iterations of the solver. Furthermore, Rioux et al. (2023) recently demonstrated that increasing enhances the convexity of the entropic GW problem, thereby improving the convergence rate of the mirror descent scheme. This provides theoretical justification for using a larger & to reduce the number of mirror descent steps. In other words, it reduces the number of calls to Sinkhorn, i.e., the number of outer iterations the solver.
We have empirically validated this behavior on DSprites with Beta + GMG. We plot the mean amount of inner and outer iterations for the first 3 epochs over 5 seeds for six different values of εo, which we provide in Figure 5 We can observe the expected scaling of increased inner and outer iterations with decreased En-
(ii) Approximation. /Zhang et al. (2033) show that when using inner products and squared Euclidean distances as costs, and for & E (0, 1,. the approximation error scales as:
IGW
GW(p.)) 4)1 S e clog(1/e),
where the constants in X depend on dx and dy, that is, the dimensions of the support of p and q. as well as their fourth-order moments.
Given this trade-off, the practical goal is to select an E value that is sufficiently large to ensure fast convergence while avoiding any degradation in performance. Across all our experiments, we found that e = 0.1 struck the right balance. We validate this observation in the next paragraph.
Algorithm 2 Entropic Gromov- Wasserstein solver (Peyré et al. 2016). (Scetbon et al. 2022 Alg. 2).
1: Require: samples x...... Xn, ~ pi yI..... Yn ~ g: cost functions cr, ey: entropic regulariza-
tion scale Eo (default = 0.1), statistic operator cost matrix stat (default mean).
2: Cx + [xx(x;, (x
> usually (n?dxx
3: Cy + [cy(y,, yr))
b usually (nddyy
4: E stastat Cx) stat(Cy)
usually O(n²)
5: P, + 1,,1,
D n
6: while converged do
7: C.+1 + C2
> n or ndx + +dy)
8: P.+1 < SINKIIORN(CHII..
O(n²)
9: end while
10: Compute GWe(Pn: qn) from P, using Eq. (EGWP)
> n or n (dx + dy)
11: return GW&(Pno 9n)
We investigate the effect of the entropic regularization strength e used to compute GMG on the disentanglement performances. The results are presented in Fig. 5 and show that performance is robust to the choice of entropic regularization scale Eu- We observed this both with respect to a setting, where we see major improvements from the GMG (DSprites with Beta as well one where we only observe minor improvements (DSprites with BetaTCVAB). This validates our choice to use a single reasonable value for all our experiments, namely Eo = 0.1. This robustness with respect to the entropic strength also observed in recent work (Piran et al., 2024) proposing a similar gap regularization, based on the entropic multi- marginl OT problem. Fig. 3) for experiments highlighting this robustness. Our intuition is as
Effect of E in disentanglement.
See (Piran al. 2024
28

---

Published as a conference paper at ICLR 2025
SmalINORB
Runtime Analysis on SmalINORB VAB
BebvVE GMG
Outer lecations
31
BELEE DMG
Berations
90
RVAE
DCTDD
Iterations
Banc BeasVA
20
0.05
011
0.3 0.3
0.5
1.0
0.68
0.1 0.2 0.3 0.5
Figure 5: Analysis of the effect of the entropic regularization scale En on the disentanglement performances when learning with the GMG. We use the SmaIINORB dataset and consider the 3-VAE (in blue) and 3-TCVAE (in orange) settings. The GMG is applied to the decoder de with cost functions cx cy cos-sim(-). corresponding to the setting that provides better disentanglement. We compute the GMG using Algorithm|| with stat mne We investigate the effect of the entropic regularization scale Eg by testing five other values of Eo besides the one used in all other experiments in the paper (namely, Eo 0.1). The values tested are Eo {0.05, 0.2, 0.3. 0.5, 1}. We also include the baseline result without using the GMG (dashed line) as comparison. Additionally, we provide runtime analysis with respect to both the inner and outer iteration of the GW-solver.
Table 5: Hyperparameter grid search for Shapes3D 128 x 128.
Method
Parameter
Values
3-(ТC)VAE
[10, 16] [0.1, 1, 10]
+ GMG
À
V-AE)VVAE
B
[2, 4]
+ GMG
X
[0.1, 1, 10]
+ HHS
7
[1, 10]
follows: when using Sinkhorn (in the case of M3G) GW/quadratic OT (in work) to compute a training loss, acts as a sharpness parameter, emphasizing certain pairs of points more strongly. While the loss value changes with e, the optimization of network variables on top of this loss appears to be largely unaffected by the sharpness introduced by e. However, when the output of Sinkhorn or GW is used directly for predictions learning flows, for example, in Monge maps (Pooladian and
Niles- Weed 2021 Kassraie et al. 2024) or Gromov-Mong maps (Klein et al. 2024), the entropic regularization strength E has much stronger influence and requires careful tuning.
D.4 SCALING TO HIGHER IMAGE RESOLUTIONS
As detailed in $3.2 the computation of the GMG scales linearly with the data dimension, enabling our method to handle high-dimensions settings effectively. To demonstrate this scalability, we benchmark our approach on the Shapes3D dataset upsealed to 128 128 resolution. To accommodate the increased image resolution, we extend the Decoder by adding one additional layer. Using our best-performing configuration from the 64 x 64 experimentsCos costs and the GMG applied to the decoder dewe conduct focused hyperparameter search as described in Table 5 The results, summarized in Table compare the four base models Beta VAE (+ HFS), both with and without the GMG. Our findings indicate that the proposed setup scales higher resolutions while preserving the performance improvements achieved by the GMG. Finally, we visually validate these results by plotting the latent traversals for the pesf-performing configuration, as shown in Figure 6
29

---

Published as a conference paper at ICLR 2025
Table Impact of the GMG, applied with Cos as the function on the decoder dø, on disentan- glement performance using upscaled 128 x 128 Shapes3D images. Performance is evaluated using DCI-D, with the best result highlighted for each method.
With Cos costs
3-VAE
B-TCVAE 3-VAE + HFS
3-TCVAE + HFS
Shapes3D (128 x 128) (Burgess and Kim 2018)
Base
54.3 19..3
74.6 17.3
87.2 +227
88.6 1114
+ Dec-GMG
63.9 +.2
82.0 12.8
90.4 #3.7
92.2 882
D.5
LATENT TRAVERSAL VISUALIZATION
Figure 6: Latent traversal visualization for Shapes3I 128 x 128 with our best performing setup, BetaTCVAE + HFS + GMG. We select the best performing result out seeds achieving a DCI-D of 99.4. We plot four different latent dimensions while traversing them from -1.0 to 1.0. As visualized the model has clearly learned to separate wall hue, object hue, scale, and floor hue into different latent dimensions.
Figure 7: Latent traversal visualization for Shapes3D 64 x 64 with our best performing setup,
BetaTCVAE + HFS + GMG. We select the best performing result out 5 seeds achieving a DCI-D of 100.0. We plot four different latent dimensions while traversing them from -1.0 to 1.0. As visualized the model has clearly learned to separate wall huc, object huc, scale, and floor hue into different latent dimensions.
30

---

Published as a conference paper at ICLR 2025
D.6 TOWARDS DECODER-FREE DISENTANGLEMENT
Recently, works such as Burns et al. 2021 Von Kgelgen
Table 7: Disentanglement (DCI-D) with- out decoder trained with various regular-
al. 202T Eastwood al. 2023; Matthes al. 2023] Aitchison and Ganev] 2024) have s shown the possibility of disentanglement through self-supervised, contrastive learning objectives in an effort to align with the scalabil- ity of encoder-only representation learning (Chen et 2020b Zbontar et al., 2021 Bardes et al., [022] Garrido
izations on Shapes3D (Burgess and Kim
2018).
Decoder-free
8-VAE
3-TCVAB
Base
0.0 +0.0
0.0 #0.0
ct alla]]]200) 2023). However, these encoder-only approaches still require weak supervision or access to multiple views of an image to learn meaningful representations of the data samples.
L2 :11 l12
+ DST
38.2 +0.8
42.7 116
+ GMG
13.9 00.4
20.5 =0.5
As the goal of geometry preservation connects the data manifold and the latent domain through minimal dis- tortion objective and is applicable to both the encoder and decoder of VAE ($3) Table 4. we posit that its ap- plication may provide sufficient training signal to learn meaningful representations and encourage disentangle- ment, eliminating the need for reconstruction loss and decoder. Table 7 shows preliminary results on unsuper- vised decoder- free disentangled representation learning on the Shapes3ID benchmark, where the decoder and associated reconstruction objective have been removed.
ScL2 : all
11 a > 0 learnable
DST
45.6 1.2
53.5 110
+ GMG
15.2 #0.3
25.2 #0.6
Cos cos-sim(s)) (s..)
+ DST
37.0 #0.4
46.1 1.5
GMG
37.0 #0.9
38.8 #1.1
Standard approaches such as 3-VAE 3-TCVAE collapse and do not achieve measurable disentan- glement (DCI-D of 0.0). However, the inclusion of either DST or GMG significantly raises achievable disentanglement and, combined with the 3-TCVAE matching objective, can achieve DCI-D scores of up to 53.5 without needing any decoder or reconstruction loss. While these are preliminary insights, we believe they offer promise for more scalable approaches to unsupervised disentangled representation learning and potential bridges to popular and scalable self-supervised representation learning approaches. Note, that here the distortion loss significantly outperforms the GMG. This expected due to the nature of the GMG, as the distortion loss offers a more restrictive and, thus, stronger signal for learning representations, which necessary in the absence of a reconstruction objective. This highlights that while in most scenarios ($ 2. Figure the GMG is preferable over the distortion loss, there also exist settings where restrictive optimization signal is desirable.
31

---

Published as a conference paper at ICLR 2025
E PYTHON CODE FOR THE COMPUTATION OF THE GROMOV-MONGE GAP
import jax
import jax.numpy jnp
from ott gecmmeryy import costs, gecmetry
from ott vers. quadratic import gromov_wasserstein
from ott.problens. w.w..rrraaa....aaaaatii..... import quadratic_problent
def gromev_ gap from srom_ssmplees
source: jax.Array:
target: jax.Array,
cost_fn: costs.CostRP
costs.cosine (),
epsilen: float 0.1, stat_fn: Callable, usually computes the mean of the cost matrix ..kwargs,
float:
"*"Gromov Monge gap regularizer on samples.
define source and target geometries
cost_mat _est__marix_x cost_fn.all htt._thh.a._ppiisssssssser y=source)
scale_cost_x stat_fniscclee_oostt>> cost_matrix_x cost_matrix_x/ geom_xx httprr/ceeottttttca.atooorroooaa
cost_matrie m_ttrittiixyy.. cost fn. fh.aall_parrs y=target)
coale_cost.. stat_fn (cst_maiicss__mttix__)))) cost_matrix_y cost_matrix_y otgp (scalcccostt geomlyy http:y//..ootttccc..oaaatroooooo
define and solve entropic GW problem prob pmaaratii_prooleemnauuaaaaaaeeee (ggoon_xx, geom_yy)
solver gromov_ epsilon-epsilon, + kwargs
out selver (prob)
compute the distortion inn induced by the map distorticn_cant = inp. n (gecm_ xx, cost_matrix yy. ctsttmmttttttttil **2
compute optimal (entropic) gromov-monge displacement reg_gw_cost. out out.reg_gtt reg__gw_cost
ent_reg_gw_cost reg_gw_cost epsilos lon jnp. log (source)) )
compute gromov-monge gap
loss = distortion_cost ent cost
return loss jax. lax. stop gradient sccale_cost x scale cost_ y)
32


Reviews:
[{"review_id": "sWZoaR5Cxk", "reviewer": "Reviewer_Fss4", "summary": "This paper introduces a novel approach for unsupervised disentangled representation learning, aiming to balance prior alignment with the preservation of geometric features within latent spaces. Achieving disentangled representations without supervision is difficult; recent studies suggest that maintaining the geometric structure of the data distribution can promote disentanglement. However, aligning data with a prior distribution while fully preserving geometric properties often requires trade-offs. Therefore, the authors introduce the Gromov-Monge Gap (GMG), a regularizer based on optimal transport theory that learns a mapping with minimal distortion.", "questions": "- Considering the entropic regularization parameter plays a role in the optimization, how does adjusting the hyperparameter affect the computational efficiency and convergence of the GMG computation? An ablation study of varying hyperparameters would be nice.\n- How scalable is GMG to high-dimensional real world datasets for the VAE configuration?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-organized and clearly written.\n- Introducing the Gromov-Monge Gap as a regularizer provides a fresh approach to minimizing geometric feature distortion, presenting a promising advancement for unsupervised disentangled representation learning.\n- The authors offer valuable insights, particularly in showing that preserving angles through cosine similarity can enhance disentanglement more effectively than preserving distances.\n- The proof demonstrating the GMG as a weakly convex regularizer is a useful contribution for future work in this area.", "weaknesses": "- The paper provides limited discussion on the scalability of GMG for high-dimensional real world datasets.\n- There is an absence of ablation studies on key hyperparameters, such as entropic regularization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "9MxxNjR5O8", "reviewer": "Reviewer_MkTy", "summary": "This paper introduces a new method for disentangled representation learning based on quadratic optimal transport theory using Gromov-Monge maps. The authors leverage Gromov-Monge Gap (GMG) to measure how well a mapping preserves geometric features while transporting between distributions with minimal distortion. The approach aims to effectively combine geometric constraints with prior matching by finding mappings that optimally balance geometric preservation and distribution alignment. The method is validated on four standard disentanglement benchmarks (Shapes3D, DSprites, SmallNORB, Cars3D) showing consistent improvements over existing disentanglement approaches.", "questions": "1. Why does angle preservation work better than distance preservation? Is this finding specific to the datasets tested?\n\n2. How sensitive is the method to the entropic regularization parameter $\\epsilon_0$?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper provides a fresh perspective on combining geometric constraints with prior matching in disentanglement learning. Different from previous works that use direct penalties for geometric preservation, the authors propose using Gromov-Monge maps to find mappings that optimally preserve geometric features while aligning distributions. The theoretical analysis proves weak convexity of the GMG regularizer and provides precise characterization of convexity constants.\n\n2. The authors have conducted comprehensive experiments across multiple benchmarks and architectural choices. Through thorough ablation studies, they demonstrate that angle preservation outperforms distance preservation, decoder regularization is more effective than encoder regularization, and the GMG shows better gradient stability compared to distortion-based approaches. Additionally, they show the method can be integrated with existing approaches like β-VAE and β-TCVAE to enhance their performance.", "weaknesses": "1. It remains unclear to me why the Gromov-Monge map framework is the optimal choice for combining geometric preservation with distribution alignment. While the authors show empirically that GMG outperforms direct distortion penalties, they don't fully explain why measuring suboptimality in the Gromov-Monge problem provides a better learning signal than other potential optimal transport formulations or geometric metrics.\n\n2. My another concern is about the practicality of solving Gromov-Wasserstein problems during training. Although the authors provide computational complexity analysis and propose using entropic regularization with scaled costs, they don't thoroughly investigate how different solver parameters affect the trade-off between computational efficiency and accuracy of the approximation. The sensitivity analysis of the entropic regularization parameter $\\epsilon_0$ and its impact on training dynamics needs more exploration.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "8h3X9rb4yb", "reviewer": "Reviewer_zj5w", "summary": "This paper studies learning disentangled representations in an unsupervised setting by introducing Gromov-Wasserstein formulation into the distortion framework. They propose GMG quantifying the difference between the distortion and the Gromov-Wasserstein distance, meanwhile, they show that GMG is weakly convex. They validate the effectiveness of GMG in four 3D datasets.", "questions": "1. Have you applied the method on datasets other than 3D dataset? For example, MNIST. \n\n2. Have you tested the method on noisy point clouds? Or any other noisy dataset? \n\n3. The whole framework is very computational expensive, do you think how it can be applied to large-scale setting? Have you thought about acceleration?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The intuition of this idea is good, it tries to minimize the difference between the previous Distortion and Gromov-Wasserstein distance so that the mapping between source and target should preserve the full features with minimal distortion. \n2. They show that the GMG and its counterpart are weakly convex functions. \n3. The experiments in 3D Shapes dataset and the clustering illustration sound persuasive.", "weaknesses": "1. The computation of Gromov-Wasserstein distance is very expensive even when entropy regularizer is employed, the time complexity is still around O(N^5). This makes the GMG not practical when involving the method in a large-scale data setting. \n2. Limited experiments: the datasets in the experiments are all about 3D shapes. How about 2D shapes? I would like to see how it can apply to the MNIST, another simple 2D digit. \n3. I believe that the robustness should also be important part but this work seems to ignore it. What if the data has out-of-distritbution points? In fact, DST is not robust in the noisy data setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ilj1lQZ5LA", "reviewer": "Reviewer_5P64", "summary": "This paper proposes to learn disentangled representation from the perspective of preserving geometric features. They propose to learn a pair of isometric encoder and decoder, hoping these inductive biases can benefit disentanglement. Based on this motivation, they propose to leverage Gromov-Monge mappings to enforce such an optimal transport (OT) mapping.", "questions": "Please refer to the weaknesses.\n\nMy main confusion comes from the logical connection of this paper from isometry to disentanglement to OT. Please try to address this issue in the rebuttal.", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "1. The empirical performance is good, indicating that adding this neural GM gap can improve the disentanglement score for many architectures.", "weaknesses": "1. The motivation and logical connection is not clear. There are three concerned concepts in the paper: isometry, OT, and disentanglement. The authors claim that adding geometrical constraints (isometry) can help disentanglement, but this connection is not explained well in the paper, which makes it hard to understand how disentanglement connects isometry. Even if there is such a connection, it is not clear why OT is required for isometry. You can just do an isometry mapping. I would expect the authors to explain these two connections more clearly in the revision.\n\n2. The experiments are not set up well to support the main contribution of the paper. If the authors would like to show the benefit of isometry, the authors can follow homeomorphic VAE [1] to show the isometry can hep data on different groups/manifolds. If the authors would like to show the benefit of OT, there is Wasserstein auto-encoder. If the authors want to show the benefit of OT in disentanglement, there are [2] and [3] which learn an OT mapping to model disentangled sequential transformations. I do not see a direct connection between isometric OT mapping and disentanglement.\n\n3. How does the method compare with the vanilla differentiable Sinkhorn algorithm [4]? Will there be any theoretical guarantees? Beyond the vanilla OT problem, there exist other OT solver like entropic OT and unbalanced OT. Do these techniques apply to disentanglement?\n\n4. In a disentanglement paper, it is rare to not see figures about latent traversal which show the exact results of disentanglement. \n\n5. How does the method work in more scaled settings, i.e., high-resolution disentanglement datasets like Falcol3D and Issac3D [5]?\n\n>[1] Explorations in Homeomorphic Variational Auto-Encoding. ICML 2019.\n>\n>[2] Flow Factorized Representation Learning. NeurIPS 2023.\n>\n>[3] Unsupervised Representation Learning from Sparse Transformation Analysis. Arxiv 2024.\n>\n>[4] Sinkhorn distances: Lightspeed computation of optimal transport. NeurIPS 2013.\n>\n>[5] Semi-Supervised StyleGAN for Disentanglement Learning. ICML 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)