MMD-Regularized Unbalanced Optimal Transport
Authors: Piyushi Manupriya
SakethaNath Jagarlapudi
Pratik Jawanpuria
Venue: Accepted by TMLR ()

Abstract: We study the unbalanced optimal transport (UOT) problem, where the marginal constraints are enforced using Maximum Mean Discrepancy (MMD) regularization. Our work is motivated by the observation that the literature on UOT is focused on regularization based on $\phi$-divergence (e.g., KL divergence). Despite the popularity of MMD, its role as a regularizer in the context of UOT seems less understood. We begin by deriving a specific dual of MMD-regularized UOT (MMD-UOT), which helps us prove several useful properties. One interesting outcome of this duality result is that MMD-UOT induces novel metrics, which not only lift the ground metric like the Wasserstein but are also sample-wise efficient to estimate like the MMD. Further, for real-world applications involving non-discrete measures, we present an estimator for the transport plan that is supported only on the given ($m$) samples. Under certain conditions, we prove that the estimation error with this finitely-supported transport plan is also $\mathcal{O}(1/\sqrt{m})$. As far as we know, such error bounds that are free from the curse of dimensionality are not known for $\phi$-divergence regularized UOT. Finally, we discuss how the proposed estimator can be computed efficiently using accelerated gradient descent. Our experiments show that MMD-UOT consistently outperforms popular baselines, including KL-regularized UOT and MMD, in diverse machine learning applications.

Published in Transactions on Machine Learning Research (01/2024)
MMD-Regularized Unbalanced Optimal Transport
Piyushi Manupriya
cs15mm220p1100000kkhh....
Department of Computer Science and Engineering, HI Hyderabad, INDIA.
J. SakethaNath Department of Computer Science and Engineering, 11 Hyderabad, INDIA.
saketha@cse.
Pratik Jawanpuria Microsoft, INDIA.
Reviewed on OpenReview: https: openreview. net/ forum? id= eN99,UU3h1
Abstract
We study the unbalanced optimal transport (UOT) problem, where the marginal constraints enforced using Maximum Mean Discrepancy (MMD) regularization. Our work is moti- vated by the observation that the literature on UOT is focused regularization based on g-divergence (e.g., KL divergence). Despite the popularity of MMD, its role as regular- izer in the context of UOT seems less understood. We begin by deriving a specific dual of MMD-regularized UOT (MMD-UOT), which helps us prove several useful properties. One interesting outcome of this duality result that MMD-UOT induces novel metrics, which not only lift the ground metric like the Wasserstein but also sample-wise efficient timate like the MMD. Further, for real-world applications involving non-discrete measures, present an estimator for the transport plan that is supported only on the given (m) samples. Under certain conditions, we prove that the estimation error with this finitely- supported transport plan is also O11//m). As far as we know, such error bounds that are free from the curse of dimensionality are not known for o-divergence regularized UOT. Finally, we discuss how the proposed estimator can be computed efficiently using acceler- ated gradient descent. Our experiments show that MMD-UOT consistently outperforms popular basclines, including KT-regularized UOT and MMD, in diverse machine learning applications.
1 Introduction
Optimal transport (OT) is a popular tool for comparing probability measures while incorporating geometry over their support. OT has witnessed lot of success in machine learning applications (Peyré & Cuturi 2019 where distributions play central role. The Kantorovich's 's formulation for OT aims to find an optimal plan for the transport of mass between the source and the target distributions that incurs the least expected cost of transportation. While classical OT strictly enforces the marginals of the transport plan to be the source and target, one would want to relax this constraint when the measures are noisy (Frogner et al. 2015) or when the source and target are un-normalized (Chizat) 2017 Licro al.) 2018). Unbalanced optimal transport (UOT) (Chizal 2017), variant of OT, employed in such cases, which performs regularization--asse noft-matching of the transport plan's marginals with the source and the target distribution..
Unbalanced optimal transport with Kullback Leibler (KL) divergence and, in general, with d- divergence (Csiszar 1967) based regularization is well-explored in literature (Liero et al. 2016 2018). Entropy regularized UOT with KL divergence (Chizat et al. 2017 2018) has been employed in applica- tions such as domain adaptation (Fatras et al.) 2021), natural language processing (Chen et al.] 2020b).
and computer vision De Plaen al. (2023). Existing works (Piccoli & Rossi, 2014 [016 Hanin 1992
Georgiou et al. 2009 have also studied total variation (oiiation(VV-regaallaiiaaaaa (TV)-reggilarizationn-aass UOT formulations. While

---

Published in Transactions on Machine Learning Research (01/2024)
MMD- methods have been popularly employed in several machine learning (ML) applications (Gretton et al.
Li et al. 2017 2021 Nguyen et al. 2021). the applicability of MMDD-baed regularization for COT is not well-understood. To the best of our knowledge, interesting questions like the following, have not been answered in prior works:
2012
Will MMD regularization for UOT also lead to novel metrics over measures, analogous to the ones obtained with the KL divergence (Liero et al. 2018) or the TV distance (Piccoli & Rossi 2014]?
What will be the statistical estimation properties of these?
How can such MMD regularized UOT metrics be estimated in practice such that they are suitable for large-scale applications?
In order to bridge this gap, we study MMD-based regularization for matching the marginals of the transport plan in the UOT formulation (henceforth termed MMD-UOT).
We first derive a specific dual of the MMD-UOT formulation (Theorem 4.1 which helps further analyze its properties. One interesting consequence of this duality result is that the optimal objective of MMD-UOT is a valid distance between the source and target measures (Corollary 4.2). whenever the transport cost is valid (ground) metric over the data points. Popularly, this is known as the phenomenon of lifting metrics to measures. This result is significant as it shows that reMD-regularizztion in UOT can parallel the metricity- preservation that happens with KL-regularization (Liero et al. 2018) and TV-regularization (Piccoli & Rossi, 2014). Furthermore, our duality result shows that this induced metric is a novel metric belonging to the family of integral probability metrics (TPMs) with a genorating set that is the interscction of the generating sets of MMD and the Kantorovich- -asoiiicc--sssstteeennn metric. Because of this important relation, the proposed distance is always smaller than the MMD distance, and hence, estimating MMD-UOT from samples is least efficient as that with MMD (Corollary 4.6). This is interesting minimax estimation rates for MMD can be completely dimension-free. As far as we know, there are no such results that show that estimation with KTegaVeregullriize UOT can be as efficient sample-wise. Thus, the proposed meirics not only lift the ground metrics to measures, like the Wasserstein, but also are sample-wiss wise efficient to estimate, like MMD.
However, like any formulation of optimal transport problems, the computation of MMD-UOT involves op- timization over all possible joint measures. This may be challenging, especially when the measures are continuous. Hence, we present a convex program-based estimator, which only involves search over joints supported at the samples. We prove that the proposed estimator is statistically consistent and converges to MMD-UOT between the true measures at a rate O (m=+), where m is the mumber of samples. Such efficient estimators are particularly useful in machine learning applications, where typically only samples from the underlying mcasures are available. Such applications include hypothesis testing, domain adaptation, and model interpolation, to name a few. In contrast, the minimax estimation rale for the Wasserstein distance is
itself O (m= (m=)), where is the dimonsionality of the samples (Niles Wood & Rigollet 2019:. That is, even if a search over all possible joints is performed, estimating Wasserstein may be challenging. Since MMD-UOT can approximate Wasserstein arbitrarily closely (as the regularization hyperparameter goes oo), our result can also be understood as a way of alleviating the curse of dimensionality problem in Wasserstein. We summarize the comparison between MMD-UOT and relevant OT variants in Table [
Finally, our result of MMD-UOT being a metric facilitates its application whenever the metric properties of OT are desired, for example, while computing the barycenter-based interpolation for single-cell RNA soquencing ( [Tong et al.) 2020). Accordingly, we also present a finieedimeensinnal convex- sarex-programm-basee estimator for the barycenter with MMD-UOT. We prove that this estimator is also consistent with an efficient sample complexity. We discuss how the formulations for estimating MMD-UOT (and barycenter) can be solved efficiently using accelerated (projected) gradient descent. This solver helps us scale well to large datasets. We empirically show the utility of MMD-UOT in several applications including two- sample hepothesis testing, single-cell RNA sequencing, domain adaptation, and prompt learning for few- shot classification. In particular, we observe that MMD-UOT outperforms popular baselines such as KL- regularized UOT and MMD in our experiments.
We summarize our main contributions below:
2

---

Published in Transactions on Machine Learning Research (01/2024)
Table 1: Summarizing interesting properties of MMD and several OT OT /UO approaches. <OT (Cuturi 2013 and cKL-UOT (Chizat 2017) denote the cntropy-regularized scalable variants OT and KL-UOT Liero
ee al. 2018), respectively. MMD and the proposed MMD-UOT are shown with characteristic kernels. By
"hhitte-araaeeeen parameterization bounds' mean results similar to Theorem 4.10
Property
MMD
OT
eTT
TV-UOT KI-UOT
eKI-UOT MMD-UOT
Metricity
X
X
Lifting of ground metric
X
V
V
V
EXP
No curse of dimensionality
X
X
X
X
Finite- paniteeparammtizaaion bounds
N/A
X
X
X
Dual of MMD-UOT and its analysis. We prove that MMD-UOT induces novel metrics that not only lift ground metrics like the Wasserstein but also are sample-wise efficient to estimate like the MMD.
Finito-dimensional convex-program-Dee based estimators for MMD-UOT and the corresponding barycen- ler. We prove that the estimators are both statistically and computationally efficient.
We illustrate the efficacy of MMD-UOT several real-world applications. Empirically, we observe that MMD-UOT consistently outperforms popular basclinc approaches.
We present proofs for all our theory results in Appendix B As side-remark, we note that most of our results not only hold for MMD-UOT but also for a UOT formulation where a general IPM replaces MMD. Proofs the appendix are hence written for general TPM-based regularization and then specialized to the case when the IPM is MMD. This generalization to IPMs may itself be of independent interest.
2 Preliminaries
Notations. Let be a set (domain) that forms a compact Hausdorff space. Let R (X), R(X) denote the of all non-negative, signed (finite) Radon measures defined over X; while the set of all probability measures is denoted by (X)). For mcasure the product space, E R+ x X), let Tı, n2 denote the first and second marginals, respectively (i.e., they are the push-forwards under the canonical projection maps onto X). Let C(X), C(X) denote the set of all real-valued measurable functions and all real-valued continuous functions, respectively, over X.
Integral Probability Metric (IPM): Given a set G C L(X), the integral probability metric (IPM) (Muller
1997 Sriperumbudur et al. 2009 Agrawal & Horel 2020) associated with G, defined by:
/ dsu o f dto & su. to E R (X).
(s to) = max feg
(1)
G is called the generating set of the IPM, 7g-
Maximum Mean Discrepancy (MMD) Let be characteristic kernel (Sriperumbudur et al. 2011 over the domain X, let Il denote the norm of f in the canonical reproducing kernel Hilbert space (RKHS).
HR, corresponding to k. MMD, is the IPM associated with the generating set: G. = ( E Hel ||/lk < 1).
Using a characteristic kernel k, MMD metric between 80, to E RR''') is defined as:
MMD (50. to) = max 1 f ds0 - -x dtol FEDx 11u (so) M4 (lv) ll
(2)
where μk (s) (x) ds(a), the kernel mean embedding of s (Muandet et al. 2017 dy is the canonical feature map of k. A kernel & called a characteristic kernel the map μs injective. MMD can be computed analytically using evaluations of the kernel k. MMD, is metric when the kernel & is characteristic. A continuous positive-defnitt kernel k X is called c-universal if the RKHS He is dense in C(X) w.r.t. the
3

---

Published in Transactions on Machine Learning Research (01/2024)
sup-norm, i.c., for every function 9 C(X) and all >>> 0, there exists an Hy such that || glocc c. Universal kernels are also characteristic. Gaussian kernel (RBF kernel) is an example of a universal kernel over the continuous domain. Dirac delta kernel is an example of a universal kernel over the discrete domain.
Optimal Transport (OT) Optimal transport provides a tool to compare distributions while incorporating the underlying geometry of their support points. Given a cost function, e: X X H R, and two probability measures 80 (2), (2),, the p- D-Wsssesseein Kantorovich OT formulation is given by:
min (so to) (ER](xx))
m di, s.t. π1 80, H2 to.
/
(3)
where p 2 1. An optimal solution of (3) is called an optimal transport plan. Whenever the cost is metric, d, over X x X (ground metric), W delines metric over measures, known as the p- Wasserstein metric, over Ri (2) x R (X).
Kantorovich metric (K.)) Kantorovich metric also belongs to the family of integral probability metrics
$1j. where e is a metric over X x X. The Kantorovich- duality result shows that the 1-Wasserstein metric is the same as the Kantorovich metric when restricted to probability measures (refer for e.g. (5.11) in Villani (2009)): )
associated with the generating set We = ʃ X R max If(x)-/(w)
=EX+vEX
c(a.v)
1 dso
. ddo dto
Wi(su, to)
min (ER((XXX)
c da,
max FEG
x
(C(ss,
s.t. 11 = 30: = to
where So, to E RT (X).
3 Related Work
Given the source and target measures, 8p R+ and to E R+ respectively, the unbalanced opti- mal transport (UOT) approach (Liero et al.) 2018 Chizat et al. 2018 aims to learn the transport plan by replacing the mass conservation marginal constraints (enforced strictly in 'balanced' OT setting) by a soft regularization/ potplarization/penlllie penalization on the marginals. KL-divergence and, in general, ø-divergence (Csiszar 1967), (Sriperumbudur et al. 2009) based regularizations have been most popularly studied in UOT setting.
The d-divergence regularized UOT formulation may be written as (Frogner el al. 2015), (Chizat 20170:
/ dl + (TI,)))) + (D2(
min " (xxxx
(4)
where c is the ground cost metric and D+(+.) denotes the d-divergence Cssisaar 1967 Sriperunbudur
et al.) 2009) between two measures. Since in UOT settings, the measures Sp, to may be un-normalized,
Tollowing (Chizat 2017: Liero et al., 2018) the transport plan is also allowed to be un-normalized. UOT
with KL-divergence- based regularization induces the so-called Gaussian Dellinger-Kantorovioo metric (Liero
et al. 2018) between the mcasures whenever 0 < x S and the ground cost с is the squared-Euclidenn
distance. Similar to the balanced OT setup (Cuturi 2013 an additional entropy regularization in KL-UOT
formulation facilitates Sinkhorn itcration (Knight, 2008) based efficient solver for KI-UOT (Chizat et al.
2017) and has been popularly employed in several machine learning applications (Tatras et al. 2021
Chen
et al. 2020b Arase et al. 2023 De Plaen et al. 2023
Total Variation (TV) distance is another popular metric between measures and the only common member of the @-divergence family and the IPM family. UOT formulation with TV regularization (denoted by ||TT) has been studied in (Piccoli & Rossi 2014h:
f di + A/T1 solTV + A/m2 tolrv-
min #ER (Xxx)
(5)
UOT with TV-divergence- based regularization induces the so-called Generalized Wasserstein metric (Piccoli & Rossi 2014] between the measures whenever X > O and the ground cost is a valid metric. As far
4

---

Published in Transactions on Machine Learning Research (01/2024)
we know, none of the existing works study the sample complexity of estimating these metrics from samples. More importantly, algorithms for solving (5) with empirical measures that computationally scale well to MT. applications seem to be absent in the literature.
Besides the family of g-divergences, the family of integral probability metrics is popularly used for comparing measures. An important member of the IPM family is the MMD metric, which also incorporates the geometry over supports through the underlying kernel. Due to its attractive statistical properties (Gretton et al. 20006)
MMD has been successfully applied in diverse set of applications including hypothesis testing (Gretton
et al. 2012 generative modelling (Li et al. 2017), self-supervised learning (Li et al. 2021 etc.
Recently, (Nath & Jawanpuria, 2020) explored learning the transport plan's kernel mean embeddings in the balanced OT selup. They proposed learning the kernel mean embedding of a joint distribution with the least expected cost and whose marginal embeddings are close to the given-sample-basee based estimates of the marginal embeddings. As kernel mcan embedding induces MMD distance, MMD-basod regularization features in the balanced OT formulation of (Nath & Jawanpurial 2020) as means to control overfitting. To ensure that valid conditional embeddings are obtained from the learned joint embeddings, (Nath & Jawanpuria, 2020) required additional feasibility constraints that restrict their solvers in scaling well to machine learning applications. We also note that (Nath & Jawanpuria 2020) neither analyze the dual of their formulation nor study its metric-related properties and their sample complexity result of O(m ) does not apply to our MMD-UOT estimator as their formulation is different from the proposed MMD-UOT formulation (6
In contrast, we bypass the issues related the validity of conditional embeddings as our formulation involves directly learning the transport plan and avoids kernel mean embedding of the transport plan. We perform detailed study of MMD regularization for UOT, which includes analyzing its dual and proving metric properties that are crucial for optimal transport formulations. To the best of our knowledge, the metricity of MMD-regularized UOT formulations has not been studied previously. The proposed algorithm scales well to large-scale machine learning applications. While we also obtain (m estimation error rate, we require a different proof strategy than (Nath & Jawanpuria| 2020). Finally, as discussed in Appendix B most of our theoretical results apply to a general IPM-regularized UOT formulation and are not. limited to the MMMD-regularized UOT formulation. This generalization does not hold for (Nath & Jawanpuria 2020
Wasserstein auto-encoders (WAE) also employ MMD for regularization. However, there are some important differences. The regularization in WAFs is only performed for one of the marginals, and the other marginal is matched exactly. This not only breaks the symmetry (and hence the metric properties) but also brings back the curse of dimensionality in estimation (for the same reasons as with unregularized OT). Further, their work does not allempt to study any theoretical properties with MMD regularization and merely employs it as a practical tool for matching marginals. Our goal is to theoretically study the metric and estimation properties with MMD regularization. We present more details in Appendix [.18
We end this section by noting key differences between MMD and OT--based approaches (including MMD- UOT). A distinguishing feature of OT-based approaches is the phenomenon of lifting the ground-metric geometry to that over distributions. One such result is visualized in Figure ((b), where the MMD-based- interpolate of the two unimodal distributions comes out to be bimodal. This is because MMD's interpolation is the (literal) average of the source and the target densities, irrespective of the kernel. This has been well-cstablished in the literature (Bottou et al.] 2017). On the other hand, OT-based approaches obtain unimodal barycenter. This is a 'geometric" interpolation that captures the characteristic aspects of the source and the target distributions. Another feature of OT-based methods is that we obtain a transport plan between the source and the target points which can be used for various alignment-based applications, c.g., cross-lingual word mapping (Alvarez-Mern & Jaakkola, 2018 Jawanpuria et al. 2020), domain adaptation (Courty et al. 2017 Courty et al.) 2017; Gurumoorthy et al. 2021 etc. On the other hand, it is unclear how MMD can be used to align the source and target data points.
5

---

Published in Transactions on Machine Learning Research (01/2024)
4 MMD Regularization for UOT
We propose to study the following UOT formulation, where the marginal constraints are enforced using MMD regularization.
Uk, to)
= min I c + A1MMD (TI, su) + (T2, tu) =ER* (Xxx)
(6)
min ʃ dx Aillus (m) - pa (8o) ||A + X2 Aull (Λ2) - uu (to) llk- (XXx)
where pa((s) is the kernel mcan embodding of s (defined Section 2 induced by the characteristic kernel k used in the generating set 9x [ Hz |l/l S 1}, and A1 are the regularization hyper-paramelers.
We begin by presenting a key duality result.
Theorem 4.1. (Duality) Whenever c, k C(N X) and X is compact, we have that:
Uk,adddd (s0:to)=
max SEBx(Ai).9E9A (Aa) Jx f dso Jx g dt
(7)
8.1. f(z) + g(v) S clz, v) 4 Va,yE
Here, Gx(A) (g E Hk
|| S A}.
The duality result helps us to study several properties of the MMD-UOT 6 discussed in the corollaries below. The proof of Theorem 4.1 is based on an application of Sion's minimax exchange theorem (Sion 1958) and is detailed in Appendix B.1
Applications in machine learning often involve comparing distributions for which the Wasserstein metric is a popular choice. While prior works have shown metric--reserree vatiion happens under KT-regularization (Liero et al., 2018) and IV-regularization (Piccoli & Rossi, 2016),, it an open question if MMMD-regularizztion in UOT also lead to valid metrics. The following result answers this affirmatively.
Corollary 4.2. (Metricity) In addition to assumptions in Theorem 4.1 whenever e is a metric, U. belongs to the family of integral probability metrics (IPMs). Also, the generating set of this IPM is the intersection of the generating set of the Kantorovich metric and the generating set of MMD, Finally, Us is a valid norm-induced induced metric over measures whenever k is characteristic. Thus, U lifts the ground metric e to that over measures.
The proof of Corollary 4.2 is detailed in Appendix B.2 This result also reveals interesting relationships be- tween UkA, the Kantorovich metric, Xe, and the MMD metric used for regularization. This is summarizod in the following two results.
Corollary 4.3. (Interpolant) In addition to assumptions in Corollary 4.2 if the kernel is c-universal (continuous and universal), then V s0, E R+(X),, limx UR.,,A,, (0) = Kc(so, Pur- ther, if the cost metric, C, dominates the charatteristcc kernel, kt induced metric, i.e., c(x,8) Z / Vk(x,2) k(y,u) 2k(xx y) V x,y E X, then Uk.c,2,2(50, lo) AMMD& (So, la) whenever 0 < 1. Finally, when X 10, MMD- UOT interpolates between the scaled MMD and the Kantorovich metric. The nature of this interpolation is already described in terms of generating sets in Corollary
We illustrate this interpolation result in Figure [ Our proof of Corollary 4.3 presented in Appendix B.3 also shows that the Euclidean distance satisfies such dominating cost assumption when the kernel employed is the Gaussian kernel and the inputs lie on a unit-norm ball. The next result presents another relationship between the metrics in the discussion.
Corollary 4.4. Us,C,A.Asss) S min (AMMD, (s. (s,,
The proof of Corollary 4.4 is st sorwigh and is presented in Appendix B.5 This result enables us to show properties like weak metrization and sample efficiency with MMD-UOT. For a sequence 8n E R2 (X'), n 2 1, we say that 8, weakly converges s E Rt (X') (denoted as 8, s), if and only if Ex=.. [x))] for all bounded continuous functions over X. It is natural to ask when is the convergence metric over measures equivalent. to weak convergence measures. The metric then said to metrize the

---

Published in Transactions on Machine Learning Research (01/2024)
Keeeel-bused cost. (5)
À 1 o0 (b)
E (1, 96) (c) UOT metrics
Figure 1: For illustration, the generating set of Santoorviic--assseeee is depicted as triangle, and the scaled generating set of MMD is depicted as disc. The intersection represents the generating set of the IPM metric induced by MMD-UOT. (a) shows the special case when our MMD-UOT metric recovers back the sample-efficient MMD metric, (b) shows the special case when our MMD-UOT metric reduces to the Kantorovich- Wasserstein metric that lifts the ground metric mcasures, and (c) shows the resulting family of new UOT metrics which are both sample-effieient and can lift ground metries measures.
weak convergence of measures or is equivalently said to weakly metrize measures. The weak metrization properties of the Wasserstein metric and MMD are well-understood (e.g., refer to Theorem 6.9 in (Villani 2009) and Theorem 7 in (Simon-Gabriel et al. 2020 ). The weak metrization property of U. follows from the above Corollary 4.4
Corollary 4.5. (Weak Metrization) Uk,.A. metrises the weak convergence of normalized measures.
The proof is presented in Appendix B.6 We now show that the metric induced by MMD-UOT inherits the attractive statistical efficiency of the MMD metric. In typical machine learning applications, only finite samples are given from the measures. Hence, it is important to study statistically efficient metrics that alleviate the curse of dimensionality problem prevalent in OT (Niles- Weed & Rigollet 2019 Sample complexity result. with the metric induced by MMD-UOT is presented as follows.
Corollary 4.6. (Sample Complexity) Let us denote Uk.AA, defined in (6). by U. Let 8m, fm denote the empirical estimates of su, to E (X') respectively with samples. Then, US fi)) - U(so, to) at rate (aport from conslants) same as that of MMDx m, 80) + 0.
Since the sample complexity of MMD with normalized characteristic kernel is ((m=b) (Smola et al, 2007), the same will be the complexity bound for the corresponding MMD-UOT. The proof of Corollary 4.6 is presented in Appendix B.7 This is interesting because, though MMD-UOT can arbitrarily well approximate Wasserstein (as À + oo), its estimation can be far more efficient than O (m-a)) which the minimax estimation rate for the Wasserstein (Niles- Weed & Rigollet, 2011).. Here, d is the dimensionality of the samples. Further, in Lemma B4 show that even when MMD; (q 2 N) is used for regularization, the sample complexity again comes out to be O (m- We conclude this section with couple of remarks.
Remark 4.7. As a side result, we prove the following theorem (Appendia B.S that relates our MMD-UOT to the MMD-regularized Kantorovich metric. We believe this connection is interesting as it generalizes the pop- ular Samnoroowcc-Ruuiissteee duality result relating (unregularized) OT to the (unregularized) Kantorovich metric.
Theorem 4.8. In addition to the assumptions in Theorem if c is a valid metric, then
so) + (t, Uk (so, fa) min Ke((,)) A MM (.ER(()
(8)
Remark 4.9. It. is noleworthy that most of our theoretical results presented in this section not only hold with the MMD-UOT formulation g but also with general IPM-regularized UOT formulation, which discuss in Appendia D This generalization may be of independent interest for future work.
Finally, minor results on robustness and connections with spectral normalized GAN MMyyato et al. 2018 are discussed in Appendix B.16 Appendix B.17 respectively.

---

Published in Transactions on Machine Learning Research (01/2024)
4.1 Finite-Samples Saatte-Samle-baseed Estimation
As noted in Corollary 4.6 MMD-UOT can be elficiently estimated from samples of source and larget. However, one needs to solve an optimization problem over all possible joint (un-normalized) measures. This can be computationally expensive (for example, optimization over the set of all joint density functions). Hence, in this section, we propose a simple estimator where the optimization is only over the joint measures supported at sample-based points. We show that our estimator statistically consistent and that the estimation is free from the curse of dimensionality.
Let m samples be given from the source, target, So. to (*) respectively" We denote D, (ma1.". ". Timm = 1,2 as the set of samples given from 8o,40 respectively. Let 8m, Im denote the empiri- cal measures using samples D2. Let us denote the Gram-matrix of D; by Ga. Let C12 be the m x m cost matrix with entries as evaluations of the cost function over D1 x D2. Following the common practice in OT literature (Chizat et al., 2017 Cuturi 2013
2
Damodaran et al. 2018 Fatras et al. 2021 Le et al.
2021 Balaji et al.) 2020; Nath Jawanpuria 2020 Peyré & Cuturi 2019), we restrict the transport plan to be supported on the finite samples from each of the measures in order to avoid the computational issues in optimizing over all possible joint densities. More specifically, let a be the m m (parameter/variaalle) /variable) matrix with entries as Oij = ( F25) where i,j E (1.... m]. With these notations and the mentioned restricted feasibility set, Problem simplifies to the following, denoted by üm (8ms (m):
Tr (aC12) + a1 a1 - -1 G. Gn x20
+ 12 l a 1- 22 G22
min
(9)
where Tr(M) denotes the trace of matrix M, | MI, and 01. 02 are the masses of the source, target measures, so, to, respectively. Since this is a Convex Program over a finite-dimensional variable, it can be solved in a computationally efficient manner (refer Scction 4.2
However, as the transport plan is now supported on the given samples alone, Corollary 4.6 does not apply. The following result shows that our estimator g is consistent, and the estimation error decays at a favourable rate.
Theorem 4.10. (Consistency of the proposed estimator) Let denote Uk. defined in 6 by U. Assume the domain X is compact, ground cost is continuous, E C(X' x X), and the kernel k is c-universal, normalized. Let the source measure (so), the target measure (to), as well as the corresponding
MMD-UOT transport plan be absolutely continuous. Also assume so(z), tu(x) > >AA x E X. Then, we have w.h.p. and any (arbitrarily small) e > that [im fm) - -((s,, to) SO Ar glc) g(e) minue ||ll| s.t. ||I clloc S e, and is the mass of the optimal MMD-UOT transport plan. Further, if belongs to Ht & Ha, then w.h.p. pitm (8mm ((00 So (1t
Vmm m + +o). Here,
Vm
We discuss the proof of the above theorem in Appendix B.9 Because h is universal, g(t) < 00 V > 0. The consistency of our estimator as m + oo can be realized, if, for example, one employs the scheme A A2 = (m4)/4) and 0 at slow enough rate such that g(e) + 0. In Appendix B.9.1 we show that even if e decays as fast as O ( (m)4))) then g(e) blows-up atmost as O (mm13)) Hence, overall, the estimation error still decays as O ( (7))). To the best of our knowledge, such consistency results have not been studied in the context of KT.regularized UOT.
4.2 Computational Aspects
Problem 9 is an instance a convex program and can be solved using the mirror descent algorithm detailed in Appendix 1.10 In the following, we propose to solve an equivalent optimization problem which helps us leverage faster solvers for MMD-UOT:
min
Tr (aCT)) + +11 a1
111--11 m
21
1 Gn +
+1
a
02
(10)
a a20RRm00
171
G22
I Note that this challenge inherent to OT (and all its variants). is not a consequence of choice MMD regularization. 2he no. samples from source and target. need not the same, in general.
8

---

Published in Transactions on Machine Learning Research (01/2024)
Algorithm Accelerated Projected Gradient Descent for solving Problem (10)
Require: Lipschitz constant L, initial oo 2 E Rmxm.
f(a) Tr (aC1) + A || = m1ll,,, + 22 11 =2 Gz
71 1.
//
= O.
while not converged do
a; = Projectzo (y:= 1
L f())).
7++1 1+vi-
2 2-1
M++ = +
(as a1-1).
N441
= i 1.
end while
return ag.
The equivalence between ( and 110p follows from standard arguments and is detailed in Appendix B.11
Our next result shows that the objective in 10 is L-smooth (proof provided in Appendix B.12 Lemma
4.11. The objective in Problem
10
L-smooth
with
L
(/(Am)))]GGnl (A2m) G 22
( 1 G22 mm).
The above result enables us to use the accelerated projected gradient descent (APGD) algorithm (Nesterov 2003 Beck & Teboulle 2009) with fixed step-size 1/L for solving 410p. The detailed steps are pre-
sented in Algorithm The overall computation cost for solving MMD-UOT 110 O O(*2), where is the optimality gap. In Section 5 we empirically observe that the APGD- solver for MMD-UOT is indeed
computationally efficient.
4.3 Barycenter
A related problem is that of barycenter interpolation of measures (Agueh & Carlier 2011 which has interesting applications (Solomon et al. 2014 2015 Gramfort et al. 2015). Given measures s1,..., 8n with total masses o,....,o, on respectively, and interpolation weights PI defined as the solution of B(51.-.. sn) = min,ER=(x)) Ei=1 PI....;Pas the barycenter s R+ is
typical applications, only sample sets, Di, from 8; are available instead of 8; themselves. Let us denote the correspouding empirical measures by ,1,...,8n One way to estimate the barycenter to consider B(s,..... ,,,). However, this may be computationally challenging to optimize, especially when the measures involved are continuous. So propose ostimating the barycenter with the restriction that the transport plan n² corresponding to Uk,n.A1,A2 (8,, is supported on D, x 1 And, let 2 0 E Rm denote the corresponding probabilities. Following (Cuturi & Doucet 2014), we also assume that the barycenter, is supported on D, Let us denote the barycenter problem with this support restriction on the transport plans and the Barycenter (m........ ..)) Let G be the Gram-matrix of D and C, be the my x matrix with entries as evaluations of the cost function.
Lemma 4.12. The barycenter problem (s........ .8,) can be equivalently written as:
min 20 Ei (Ir (a,CT) + A11la my 111. della P) 111
(11)
We present the proof in Appendix B.14.1 Similar to Problem (10). the objective Problem a smooth quadratic program in each a; and jointly convex in aç's. In Appendix (B.11.2 we also present the details for solving Problem (11 using APGD as well as its statistical consistency in Appendix B.14.3

---

Published in Transactions on Machine Learning Research (01/2024)
70
1.
CRUJOOT
347
Problem APGD
0.8
20-
344
0.2
241
6.6
<KL-UOTT plan MMD-UOT plan
Barycenter
Time
(b)
(6)
Figure 2: (a) Optimal Transport plans of cKL-UOT and MMD-UOT: (b) Barycenter interpolating between Gaussian measures. For the chosen hyperparameter, the barycenters of cKL-UOT and MMD-UOT overlap and can be looked smooth approximations of the OT barycenter; (c) Objective Time plot comparing
eKL-UOT solved using the popular Sinkhorn algorithm (Chizat et al. 2017 Pham et al. 2020) and MMD- UOT 10 solved using APGD. A plot showing eKL-UOT's progress at the initial phase is given in Figure
5 Experiments
In Section 4
we examined the theoretical properties of the proposed MMD-UOT formulation. In this section, we show that MMD-UOT good practical alternative to the popular entropy-regularized eKI-UOT. We emphasize that our purpose is not benchmark saate-of-the-at performance. Our codes are publicly available at httpp://ggtaubb.cm/iiissssi//a/ MMMtsshii-MMMMMr----0T0
5.1 Synthetic Experiments
We present some synthetic experiments to visualize the quality of our solution. Please refer to Appendix C.1 for more details.
Transport Plan and Barycenter We perform synthetic experiments with the source and target Gaussian measures. We compare the OT plan of eKL-UOT and MMD-UOT in Figure 2[a). We observe that the MMD-UOT plan is sparser compared to the <KL-UOT plan. In Figure ( we visualize the barycenter interpolating between the source and target, obtained with MMD, eKI-COT and MMD-UOT. While MMD barycenter is an empirical average of the measures and hence has two modes, the geometry of measures considered in both eKL-UOT and MMD-UOT formulations. Barycenters obtained by these methods have the same number of modes (one) as in the source and the target. Moreover, they appear smoothly approximate the barycenter obtained with OT (solved using linear program).
Visualizing the Level Sets Applications like generative modeling deal with optimization over the pa- rameter (9) of the source distribution to match the target distribution. In such cases, it is desirable that the level sets of the distance function over the measures show a lesser number of stationary points that are not global optima (Bottou et al.) 2017). Similar to (Bottou et al.] 2017), we consider a model family for source distributions
F [Po [00 + 8_0) : 0 [-1,11 1 [-1, 1,,1] and fixed target distribution Q P(2.2) & F. We compute the distances between Pg and Q according to various divergences. Figure 3 presents level sets showing the set of distances (d(Po. Q) measured using MMD, Kantorovich metric, eKL-UOT, and MMD-UOT (9), respectively. While all mothods correctly identify global minima (green arrow), level sets with MMD-UOT and eKL-UOT show no local minima (encircled in red for MMD) and have a lesser number of non-optimal stationary points (marked with black arrows) compared to the Kantorovich metric Figure bb.
[-1,1] [-1, where the distance d(-,-) is
Computation Time In Figure cc), we present the objective versus time plot. The source and target measures chosen to be the same, in which case the optimal objective is O. MMD-UOT (10) solved using
10

---

Published in Transactions on Machine Learning Research (01/2024)
M4 x4
A1
-1
Figure 3: Level sets distance function between a family of source distributions and a fixed target distri- bution with the task of finding the souree distribution closest to the target distribution using (a) MMD, (b) W2, (c) eKI-UOT, and (d) MMD-UOT. While all methods correctly identify global minima (green arrows), level sets with MMD-UOT and eKL-UOT show local minima (encircled in red for MMD) and have a lesser number of non-optimal stationary points (marked with black arrows) compared to (b).
Table 2: Average Test Power (between 0 and 1; higher is better) on MNIST. MMD-UOT obtains the highest average test power at all timesteps.
N
MMD
cKL-UOT
MMD-UOT
100
0.137
0.099
0.154
200
0.258
0.197
0.333
300
0.467
0,242
0.588
400
0.656
0.324
0.762
500
0.792
0.357
0.873
1000
0.909
0.506
0.909
APGD (described in Section 4.2 gives much faster rate of decrease in objective compared to the Sinkhorn algorithm used for solving KL-UOT.
5.2 Two-Sample Hypothesis Test
Given two sets of samples (.i,....,TIn) so and (yi...., (ym) ~ to. the two-sample test aims to determine whether the two sets of samples are drawn from the same distributions, viz., to predict if so = 4-- The performance evaluation in the two-sample test relies on two types of errors. Type-I error occurs when Sp = fo, but the algorithm prodicts otherwise. Type-II error occurs when the algorithm incorrectly predicts 80 = lo. The probability of Type-I error called the significance level. The significance level can be controlled using permutation test-based setups (Ernst) 20044] Liu et al.] 2020). Algorithms are typically compared based on the empirical estimate of their test power (higher is better), defined as the probability of not making Type-II error and the average Type-I error (lower is better).
Dataset and experimental setup. Following [Liu et al. 2020 we consider the two sets of samples, one from the true MNIST (LeCun & Cortes 2010) and another from fake MNIST generated by the DCGAN (Bian) et al. 2019). The data lies in 1024 dimensions. We take an increasing number of samples (N) and compute the average test power over 100 pairs of sets for each value of N. We repeat the experiment report the average test power in Table 2 for the significance level a 0.05. By the design of the test, the average Type-I error was upper-bounded, and we noted the Type-II error in our experiment. We detail the procedure for choosing the hyperparameters and the list of chosen hyperparameters for each method in Appendix C.2
times and
Results. In Table 2 we observe that MMD-UOT obtains the highest Lest power for all values of N. The average test power of MMD-UOT is 1.5 1.4 times better than that of eKL-UOT across N. MMD-UOT also outperforms EMD and 2- Wasscrstein, which suffer from the curse of dimensionality, for all values of N. Our results match the sample efficient MMD metric's result on increasing N to 1000, but for lesser sample-size, MMD-UOT is always better than MMD.
11

---

Published in Transactions on Machine Learning Research (01/2024)
Table 3: MMD distance (lower is better) between computed barycenter and the ground truth distribution. A sigma-heuristics based RBF kernel is usod compute the MMD distance. We observe that MMD-UOT's results are closer to the ground truth than the baselines' results al all timesteps.
Timestop
MMD
eKL-UOT
MMD-UOT
t
0.375
0.391
0.334
12
0.190
0.184
0.179
t3
0.125
0.138
0.116
Avg.
0.230
0.238
0.210
5.3 Single-Cell RNA Sequencing
We empirically evaluate the quality of our barycenter the Single-cell RNA sequencing experiment. Single- cell RNA sequencing technique (scRNA-scq) helps us understand how the expression profile of the cells changes (Schiebinger al. 2019). Barycenter estimation in the OT framework offers prineipled approach to estimate the trajectory of a measure at intermediate timestep (t, < <t<t)) when we have measurements available only at & (source) and t, (target) time steps.
Dataset and experimental setup. We perform experiments on the Embryoid Body (EB) single-cell dataset (Moon et al. 2019). The dataset has samples available at five timesteps (t, with 0,...... 4), which were collected during a 25-day period development of the human embryo. Following (Tong et al. 2020), we project the data onto two-dimensional space and associate uniform measures the source and the target samples given at different. timesteps. We consider the samples timestep and 444 as the samples from the source and target measures where 0 < i S 2 and aim at estimating the measure at & timestep as their barycenter with equal interpolation weights
P2 = 0.5.
We compute the barycenters using MMD-UOT 2018) approaches. For both, simplex constraint is used to cater to the case of uniform measures. We also compare against the empirical average of the source and largel measures, which is the barycenter obtained with the MMD metric. The computed barycenter is evaluated against the measure corresponding to the ground truth samples available at the corresponding timestep. We compute the distance between the two using the MMD metric with RBF kernel (Gretton et al., 2012). The hyperparameterss are chosen based on the leave-one-out validation protocol. More details and some additional results are in Appendix CC.3
and the <KI-UOT (Chizat et al. 2018 Liero et al.
Results. Table 3 shows that MMD-UOT achioves the lowest distance from the ground truth for all the timesteps, illustrating its superior interpolation quality.
5.4 Domain Adaptation in JUMBOT framework
OT has been widely employed in domain adaptation problems CCourty et al. 2017 Courty et al. 2017 Seguy
et al. 2018 Damodaran et al. 2018). JUMBOT (Fatras et al. 2021) is a popular domain adaptation method based on eKL-UOT that outperforms OT-based baselines. JUMBOT's loss function involves a cross-entres cross-entropy term and eKT-UOT discrepancy term between the source and target distributions. We showcase the utility of MMD-UOT 11 in the JUMBOT Fatras et al. 2021 framework.
Dataset and experimental setup: We perform the domain adaptation experiment with and Digits datasets comprising of MNIST (LeCun & Cortes 2010), M-MNIST (Ganin et al. 2016). SVHN (Netzer ct al. 2011,, USPS (Hull 1994h datasots. We replace the eKI-UOT based loss with the MMD-UOT loss 10h, keeping the other experimental set-up the same as JUMBOT. We obtain JUMBOT's result with eKL-UOT with the best-reported hyperparameters (Fatras et al. 2021). Following JUMBOT, we tune hyperparameters of MMD-UOT for the Digits experiment on USPS to MNIST (U++M) domain adaptation task and use the same hyperparameters for the rest of the domain adaptation tasks on Digits. More details are in A Appendix C.4
12

---

Published in Transactions on Machine Learning Research (01/2024)
Table 4: Target domain accuracy (higher better) obtained domain adaptation experiments. Results for cKL-UOT are reproduced from the code open-sourced for JUMBOT in (Fatras al., 2021). MMD-UOT outperforms <KI-UOT in all the domain adaptation tasks considered.
Source
Target
<KL-UOT
MMD-UOT
M-MNIST
USPS
91.53
94.97
M-MNIST
MNIST
99.35
99.50
MNIST
M-MNIST
96.51
96.96
MNIST
USPS
96.51
97.01
SVHN
M-MNIST
94.26
95.35
SVHN
MNIST
98.68
98.98
SVHN
USPS
92.78
93.22
USPS
MNIST
96.76
98.53
Avg.
95.80
96.82
Results: Table reports the accuracy obtained on target datasets. We observe that MMD-UOT-BER based loss performs better than KKL-UOT-basee loss for all the domain adaptation tasks. Figure 8 (appendix), we also compare the t-SNE plot of the embeddings learned with the MMD-UOT and the eKLL-UOT-based loss functions. The clusters learned with MMD-UOT are better separated (c.g., red- and cyan-colored clusters).
5.5 More Results Domain Adaptation
In Section 5.4 we compared the proposed AMD-UOTT- -ass- loss function with the eKL-UOT based loss function in the JUMBOT framework (Fatras et al. 2021). It should be noted that JUMBOT has a ResNet- 50 backbone. Hence, in this section, also compare with popular domain adaptation baselines having ResNet-50 backbone. These include DANN (Ganin et al. 2015 CDANN-E (Long et al. 2017), DEEPJ- DOT (Damodaran et al. 2018), ALDA (Chen et al. 2020a ROT (Baaai et al. 2020j, and BombOT (Nguyen ct al. 2022p. BombOT is a recent state-of-the-arr OT based method for unsupervised domain adaptation (UDA). As in JUMBOT (Fatras et al.) 2021). BombOT also employs <T--UOT based loss function. We also include the results of the bascline ResNet-50 model, where the model is trainod on the source and is evaluated on the larget without employing any adaptation Lechniques.
Office- Home dataset: We evaluate the proposed method on the Office-Honee dataset (Venkateswara
et al.) 2017), popular for unsupervised domain adaptation. We use the backbone network of ResNet-50 following. The Office-Home dataset has 15,500 images from four domains: Artistic images (A), Clip Art (C),
Product images (P) and World (R). The dataset contains images of 65 object categories common office and home scenarios for each domain. Following (Fatras al. 2021 Nguyen et al. 2022). evaluation is done in 12 adaptation tasks. Following JUMBOT, we validate the proposed method the A-+C task and use the chosen lyperparameters for the rest of the tasks.
Table reports the target accuracies obtained by different methods. The results of the BombOT method are quoted from (Nguyen et al.] 2022), and the results of other basclines are quoted from (Fatras et al. 2021 We observe that the proposed MMD-UOT-based method achieves the best larget accuracy out of adaptation tasks.
VisDA-2017 dataset: We next consider the next domain adaptation task between the training and validation sets of the VisDA-2017 (Recht et al.) 2018) dataset. We follow the experimental setup detailed in Fatras et al. 2021). The source domain of VisDA has 152,397 synthetic images, while the target domain has 55,388 real-world images. Both the domains have 12 object categories.
Table G compares the performance of different methods. The results of the BombOT method are quoted from Nguyen et al. 2022), and the results of other baselines quoted from Fatras et al. 2021 The proposed
13

---

Published in Transactions on Machine Learning Research (01/2024)
Table 5: Target accuracies (higher better) on the Office-Home dataset in the UDA setting. The letters denote different domains: 'A' for Artistic images, 'P' for Product images, 'C' for Clip art and 'R' for Real- World images. The proposed method achieves the highest accuracy on almost all the domain adaptation tasks and achieves the best accuracy averaged across the tasks.
Method
A-C
A-PP
A-R
C-+
C-P
C-R
P-A
P-C
P-R
R-A
R-C
R-+P
Avg
RosNet-50
34.9
50.0
58.0
37.4
41.9
46.2
38.5
31.2
60.4
53.9
41.2
59.9
46.1
DANN
Ganin et al. 2015
44.3
59.8
69.8
48.0
58.3
63.0
49.7
42.7
70.6
64.0
51.7
78.3
68.3
CDAN-E
52.5
71.4
76.1
59.7
69.9
71.5
58.7
50.3
77.5
70.5
57.9
83.5
66.6
Long 2017
DEEPJDOT
30.7
68.7
74.4
59.9
63.8
68.1
53.2
46.3
73.8
66.0
54.9
78.3
63.5
Damodaran
2018
ALDA
76.4
57.4
57.3
82.5
Chen et al.
2020
ROT
47.2
71.8
76.4
58.6
68.1
70.2
56.5
45.0
75.8
69.4
52.1
80.6
64.3
[alajii ct al.
2020
<K-UUT (JUMBOT)
55.2
75.5
80.8
65.5
74.4
74.9
65.2
52.7
79.2
73.0
59.9
83.4
70.0
Fatras al.
2021
BombOT
36.2
75.2
80.5
65.8
74.6
75.4
66.2
53.2
80.0
74.2
60.1
83.3
70.4
Nguyen et
2022
Proposed
56.5
77.2
82.0
70.0
77.1
77.8
69.3
55.1
82.0
75.5
59.3
84.0
72.2
Table 6: Target accuracies (higher is better) on the VisDA-2017 dataset in the UDA setting. The proposed MMD-UOT method achieves the highest accuracy.
Dataset
CDAN-E
ALDA DEEPJDOT ROT eKL-UOT (JUMBOT) BombOT
Proposed
VisDA-2017
70.1
70.5
68.0
66.3
72.5
74.6
77.0
method achieves the best performance, improving the accuracy obtained by eKL-UOT based JUMBOT and BombOT methods by 4.5% and 2.4%, respectively.
5.6 Prompt Learning for Few-Shot Classification
The task of learning prompts (e.g. "a tall bird of [class]") for vision-language models has emerged as a promising approach to adapt large pre-trained models like CLIP (Radford et al. 2021 for downstream tasks. The similarity between prompt features (which are class-specific) and visual leatures of given image can help us classify the image. recent OT-basod prompt learning approach, PLOT (Chen et al.] 2023 obtained stale-of-the-art results the K-shot recognition task in which only K images per class are available during training. We evaluate the performance of MMD-UOT following the setup of (Chen et al.
2023
the benchmark EuroSAT (Helber et al. 2018) dataset consisting of satellite images, DTD (Cimpoi et al. 2014) dataset having images of textures and Oxford-Pets (arkki et al. 2012 dataset having images of pets.
Results With the same evaluation protocol as (Chen et al. 2023 we report the classification accuracy averaged over three sceds in Table We note that MMD-UOT- based prompt-learning achieves better results than PLOT, especially when K is less (more challenging case due to lesser training data). With the EuroSAT dataset, the improvement. is as high as 4% for challenging of K=1. More details are in Appendix C.5.
14

---

Published in Transactions on Machine Learning Research (01/2024)
Table 7: Average and standard deviation (over 3 runs) of accuracy (higher better) on the k-shot classifi- cation task, shown for different values of shots (k) in the stato-of-tho-art PLOT framework. The proposod method replaces OT with MMD-UOT in PLOT, keeping all other hyperparameters the same. The results
of PLOT are taken from their paper (Chen et al. 2023
Datasel
Method
2
8
16
PL.OT
54.05 + 5.95
64.21 + 1.90
72.36 + 2.29
78.15 + 2.65
82.23 + 0.91
EuroSAT
Proposed
58.47 + 1.37
66.0 = 0.93
71.97 = 2.21
79.03 + 1.91
83.23 + 0.24
PLOT
46.55
2.62
51.24 + 1.95
56.03 + 0.43
61.70 + 0.35
65.60
0.82
DTD
Proposed
47.27+1.46
51.0+1.71
56.40+0.73
63.17+0069
65.90 0.29
6
Conclusion
The literature on unbalanced optimal transport (UOT) has largely focused on ø-divergence- based regu- larization. Our work provides a comprehensive analysis of MMD-reguulaizattin in UOT, answering many open questions. We prove novel results on the metricity and the sample efficiency of MMD-UOT, propose consistent estimators which can be computed efficiently, and illustrate its empirical effectiveness several machine learning applications. Our theoretical and empirical contributions for MMD-UOT and its corre- sponding barycenter demonstrate the potential of HIMD-regularization in UOT as an effective alternative to diveergen based regularization. Interesting directions of future work include exploring applications of IPM-regularized UOT (Remark 4.9 and the generalization of Serboroovic--uubistteee duality (Remark 4.7)
7 Funding Disclosure and Acknowledgements
We thank Kilian Fatras for the discussions on the JUMBOT bascline, and Bharath Sriperumbudur (PSU) and G. Ramesh (IIT)) for discussions related Appendix [.9.1 We are grateful to Rudraram Siddhi Vinayaka. We also thank the anonymous reviewers for constructive feedback. PM and JSN acknowledge the support of Google PhD Fellowship and Fujitsu Limited (Japan), respectively.
References
Rohit Agrawal and Thibaut Horel. Optimal bounds between f-divergences and integral probability metrics. In ICML, 2020.
Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43()):004-924, 2011.
David Alvarez-Melis and Tommi Jaakkola. Gromov-Wasserstel alignment of word embedding spaces. Tn EMNLP, 2018.
Yuki Arase, Han Bao, and Sho Yokoi. Unbalanced optimal transport for unbalanced word alignment. ACI, 2023.
Yogesh Balaji, Rama Chellappa, and Soheil Feizi. Robust optimal transport with applications in generative modeling and domain adaptation. NewrIPS, 2020.
Amir Beck and Marc Teboulle. A fast iterative shrinkage-thresholding algorithm for linear inverse problems. SIAM Journal on Imaging Sciences, 2(1):183 202, 2009.
A. Ben- Tal and A. Nemirovski. Lectures On Modern Convex Optimization, 2021.
Yuemin Bian, Junmei Wang, Jaden Jungho Jun, and Xiang- Qun Xic. Deep convolutional generative adver- sarial network (degan) models for screening and design of small molecules targeting cannabinoid receptors. Molecular Phormaceutics, 16(11):4451--4 2019.
15

---

Published in Transactions on Machine Learning Research (01/2024)
Alberto Bictti and Julien Mairal. Group invariance, stability to deformations, and complexity of deep convolutional representations. Journal of Machine Learning Research, 20:25:1-25:44 -5:499, 2017.
Alberto Bietti, Grégoire Mialon, Dexiong Chen, and Julien Mairal. A kernel perspective for regularizing deep neural networks. In ICML, 2019.
Leon Bottou, Martin Arjovsky, David Lopez-Paz, and Maxime Oquab. Geometrical insights for implicit generative modeling. Braverman Readings in Machine Learning 2017, pp. 229 268, 2017.
Guangyi Chen, Weiran Yao, Xiangehen Song, Xinyue Li, Yongming Rao, and Kun Zhang. Prompt learning with optimal transport for vision-language models. In ICLR, 2023.
Minghao Chen, Shuai Zhao, Haifeng Liu, and Deng Cai. Adversarial-leatne loss for domain adaptation. In AAAI, 2020a.
Yimeng Chen, Yanyan Lan, Ruinbin Xiong, Liang Pang, Zhiming Ma, and Xueqi Cheng. Evaluating natural language generation via unbalanced optimal transport. In IJCAI, 2020b.
Xiuyuan Cheng and Alexander Cloninger. Classification logit two-sample testing by neural networks for differentiating near manifold densities. IEEE Transactions on Information Theory, 68:6631 08:6631-6662, 2019.
L. Chizat, G. Peyre, B. Schmitzer, and F.-X. Vialard. Unbalanced optimal transport: Dynamic and kan-
torovich formulations. Journal of Functional Analysis, 274(11):3090 3123, 2018.
Lénaïe Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xave Xavier Vialard. Scaling algorithms for unbalanced optimal transport problems. Math. Comput., 87:2563 22056,9600,, 2017.
Lenaïc Chizat. Unbalanced optimal transport: Models, numerical methods, applications. Technical report, Universite Paris sciences et lettres, 2017.
Kacper P. Chwialkowski, Aaditya Ramdas, D. Sejdinovic, and Arthur Gretton. Fast two-sample testing with analytic represenlations of probability measures. In NIPS, 2015.
M. Cimpoi, S. Maji, I. Kokkinos, S. Mohamed, and A. Vodaldi. Describing textures in the wild. In CVPR, 2014.
Samuel Cohen, Michael Arbel, and Mare Peter Deisenroth. Estimating barycenters of measures in high dimensions. arXiv preprint arXiv:2007.07100 2020.
N. Courty, R. Flamary, D. Tuia, and A. Rakotomamonjy. Optimal transport for domain adaptation. JEEE Transactions on Pattern Analysis and Machine Intelligence, 39(9):1833-1866 2017.
Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. In NIPS, 2017.
I. Csiszar. Information-type measures of difference of probability distributions and indirect observations. Studia Scientiarum Mathematicarum Hungarica, 2:299 318, 1967.
M. Cuturi. Sinkhorn distances:
Lightspeed computation of optimal transport. In NIPS, 2013.
Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In ICML, 2014.
Bharath Bhushan Damodaran, Benjamin Kellenberger, Rémi Flamary, Devis Tuia, and Nicolas Courty, DeepJDOT: Deep Joint Distribution Optimal Transport for Unsupervised Domain Adaptation. ECCV, 2018.
Henri De Plaen, Pierre-Prançois De Plaen, Johan A. K. Suykens, Mare Proesmans, Tinne Tuytelaars, and Luc Van Gool. Unbalanced optimal transport: A unified framework for object detection. In CVPR, 2023.
Michael D. Ernst. Permutation Methods: A Basis for Exact Inference, Statistical Science, 19(4)6676 685,
2004.
16

---

Published in Transactions on Machine Learning Research (01/2024)
Kilian Fatras, Thibault Séjourné, Nicolas Courty, and Rémi Flamary. Unbalanced minibatch optimal trans- port; applications to domain adaptation. In ICML, 2021.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Cham- bon, Lactitia Chapel, Adricn Corenfios, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalic T.H. Gayraud, Hicham Janati, Alain Rakotomamonijy, levgen Redko, Autoine Rolet, Antony Schulz, Vivien Seguy, Danica Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya-Polo, and Tomaso Poggio. Learning with wasserstein loss. In NIPS, 2015.
Yaroslav Ganin, E. Ustinova, Hana Ajakan, Pascal Germain, H. Larochelle, François Laviolette, Mario Marchand, and Victor S. Lempitsky. Domain-adversarial training of neural networks. In Journal of Machine Learning Research, 2015.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Comain-adversariio training of neural networks. Journal of Machine Learning Research, 17(1))2096 2030, 2016.
Tryphon T. Georgiou, Johan Karlsson, and Mir Shahrouz Takyar. Metrics for power spectra: An axiomatic approach. IEEE Transactions on Signal Processing, 57(3):859 867, 2009.
Alexandre Gramfort, Gabriel Peyré, and Marco Cuturi. Fast optimal transport averaging of neuroimaging data. In Proceedings of 24th International Conference on Information Processing in Medicall Imaging, 2015.
Arthur Gretton. A simpler condition for consistency of a kernel independence test. arXiv: Machine Learning, 2015.
Arthur Gretion, Karsten M. Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander J. Smola. A kernel method for the two-sample-problem In NIPS, 2006.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(2)):723 773, 2012.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Improved training of wasserstein gans. In NIPS, 2017.
K. Gurumoorthy, P. Jawanpuria, and B. Mishra. SPOT: A framework for selection of prototypes using optimal transport. In European Conference on Machine Learning and Knowledge Discovery in Databases (ECML PKDD), 2021.
Leonid G. Hanin. -untonorich-rubinstinn norm and its application in the theory of lipschitz spaces. Proceedings of the Americal Mathematical Sociely, volume 115, 1992.
Patrick Helber, Benjamin Bischke, Andreas Dengel, and Damian Borth. Introducing eurosat: A novel datasct and deep learning benchmark for land use and land cover classification. In IGARSS 2018-2018 IEEE International Geoscience and Remote Sensing Symposium, pp. 204-207. IEEE, 2018.
J.J. Hull. A database for handwritten text recognition research. IEEE Transactions Pallern Analysis and Machine Intelligence, 16(5):550-58) 554, 1994.
P. Jawanpuria, M. Meghwanshi, and B. Mishra. Geometry-aware domain adaptation for unsupervised align- ment of word embeddings. In Annual Meeting of the Association for Computational Linguistics, 2020.
Wittawat Jitkrittum, Zoltán Szabó, Kacper Chrialkowski, and Arthur Gretton. Interpretable distribution features with maximum testing power. In NIPS, 2016.
17

---

Published in Transactions on Machine Learning Research (01/2024)
Philip A. Knight. The makhorn-knopp algorithm: Convergence and applications. SIAM Journal on Matrix Analysis and Applications, 30(1):261 225,, 2008.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
Khang Le, Huy Nguyen, Quang M Nguyen, Tung Pham, Hung Bui, and Nhat Ho. On robust optimal transport: Computational complexity and barycenter computation. In NeurIPS, 2021.
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010.
/exdb//e//////// /mnist/.
Chun-Liang Li, Wei-Cheng Chang, Yu Cheng, Yiming Yang, and Barnabás Póczos. MMD GAN: Towards Deeper Understanding of Moment Matching Network. In NIPS, 2017.
Yazhe Ti, Roman Pogodin, Danica J. Sutherland, and Arthur Gretton. Self-supervised learning with kernel dependence maximization. In NewrIPS, 2021.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal transport in competition with reaction:
The hellinger-kantorovide distance and geodesic curves. SIAM Math. Anal., 48:2869 2911, 2016.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and new
hellinger kantorovich distance between positive measures. Inventiones mathematicae, 211(3):969 1117, 2018.
Feng Liu, Wenkai Xu, Jie Lu, Guangquan Zhang, Arthur Gretton, and Danica J. Sutherland. Learning deep kernels for non-paraettric two-sample tests. In ICML, 2020.
Mingsheng Long, Zhangjie Cao, Jianmin Wang, and Michael Jordan. Conditional adversarial domain adaptation. In NIPS, 2017.
David Lopez-Paz and Maxime Oquab. evisiting classifier two-sample tests. In ICLR, 2017.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for gener- ative adversarial networks. In ICLR, 2018.
Kevin R. Moon, David van Dijk, Zheng Wang, Scott Gigante, Daniel B. Burkhardt, William S. Chen, Kristina Yim, Antonia van den Elzen, Matthew J. Hirn, Ronald R. Coifman, Natalia B. Ivanova, Guy Wolf, and Smita Krishnaswamy. Visualizing structure and transitions for biological data exploration. Nature Biotechnology, (7(1))11482-B1 1492, 2019.
Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur, and Bernhard Schölkopf. Kernel mean bedding of distributions: A review and beyond. Foumdations and Trends@ in Machine Learning, 10(1-2):
1 1-141, 2017.
Alfred Muller. Integral probability metries and their generating classes of functions. Advances in Applied Probability, 29:429 443, 1997.
J. Saketha. Nath and Pratik Kumar Jawanpuria. Statistical optimal transport posed as learning kernel embedding. In NewrIPS, 2020.
Yurii Nesterov. Introductory lectures on conver optimization: basic course, volume 87. Springer Science Business Media, 2003.
Yuval Netzer, Tiejic Wang, Adam Coates, A. Bissacco, Bo Wu, and A. Ng. Reading digits in natural images with unsupervised feature learning. In NeurIPS, 2011.
Khai Nguyen, Dang Nguyen, Quoc Nguyen, Tung Pham, Hung Bui, Dinh Phung, Trung Lc, and Nhat Ho. On transportation of mini-batches: A hierarchical approach. In ICML, 2022.
Thanh Tang Nguyen, Sunil Gupta, and Svetha Venkatesh. Distributional reinforcement learning via moment matching. In AA 2021.
18

---

Published in Transactions on Machine Learning Research (01/2024)
Jonathan Niles- Weed and Philippe Rigollet. Estimation of Wasserstein distances in the spiked transport model. In Bernoulli, 2019.
O. M. Parkhi, A. Vedaldi, A. Zisserman, and C. V. Jawahar. Cats and dogs. In CVPR, 2012.
Gabriel Poyré and Marco Cuturi. Computational optimal transport. Foundations and Trends8 in Machine Learning, 11(5-6):355-600 607, 2019.
Khiem Pham, Khang Le, Nhat Ho, Tung Pham, and Hung Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In ICML, 2020.
Benedetto Piccoli and Francesco Rossi. Generalized wasserstein distance and its application to transport
equations with source. Archive for Rational Mechanics and Analysis, 211:335 358, 2014.
Benedetto piccoli and Francesco Rossi. On properties of the generalized wasserstein distance. Archine for Rational Mechanics and Anolysis, 222, 12 2016.
Alec Radford, Jong Wook Kim, Chris Hallacy Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual models from natural language supervision. In ICML, 2021.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do CIFAR-10 classifiers gener- alize to CIFAR-10? arXiv, 2018
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, Lia Lee, Jenny Chen, Justin Brumbaugh, Philippe Rigollet, Konrad Hochedlinger, Rudolf Jaenisch, Aviv Regev, and Eric S. Lander. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928 943.e22, 2019.
Vivien. Seguy, Bharath B. Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Targe-scale optimal transport and mapping estimation. In JCLR, 2018.
Carl-Johann Simon-Gabrie Alessandro Barp, Bernhard Schölkopf, and Lester Mackey. Metrizing weak convergence with maximum mean discrepancies. arXiv, 2020.
Maurice Sion. On general minimax theorems. Pacific Journal of Mathematics, 8(1):171 176, 1958.
Alexander J. Smola, Arthur Gretton, Le Song, and Bernhard Schölkopf. A hilbert space embedding for distributions. In ALT, 2007.
Justin Solomon, Raif Rustamov, Leonidas Guibas, and Adrian Butscher. Wasserstein propagation for semi- supervised learning. In ICML, 2014.
Justin Solomon, Fernando de Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transportation geometric domains. ACM Trans. Graph., 34(4), 2015.
L. Song. Learning via hilbert space embedding of distributions. In PhD Thesis, 2008.
Khurram Soomro, Amir Roshan Zamir, and Mubarak Shah. UCF101: datasct of 101 human actions classes from videos in the wild. CoRR, 2012.
Bharath K. Sriperumbudur, Kenji Fukumiizu, Arthur Gretton, Bernhard Schölkopf, and Gert R. G. Lanckrict. On integral probability metrics, phi-divergences and binary classification. arXiv, 2009.
Bharath K. Sriperumbudur, Kenji Fukumizu, and Gert R. G. Lanckriet. Universality, characteristic kernels and RKHS embedding of measures. Journal of Machine Learning Research, 12:2389-2411 2011.
Ilya O. Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schölkopf. Wasserstein auto-encoders. ICLR, 2018.
19

---

Published in Transactions on Machine Learning Research (01/2024)
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. TrajectoryNet: A dynamic optimal transport network for modeling cellular dynamics. In ICML, 2020.
Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan, Deep hashing network for unsupervised domain adaptation. In CVPR, 2017.
Cédric Villani. Optimal Transport: Old and New. A series Comprehensive Studies in Mathematics. Springer, 2009.
A Preliminaries
A.1 Integral Probability Metric (IPM):
Agrawal & Horel 2020) associated with G, is defined by: Given a set G [(X), the integral probability metric (IPM) (Muller 1997
Sriperumbudur et al.
2009
J dsp f f dio A 80. to E (2).
(8(80, (o) = max
(12)
ʃE9
G is called the generating set of the IPM, 7g.
In order that the IPM metrizes weak convergence, assume the following Muuler 1997 Assumption A.1. C C(X) and is compact.
Since the IPM generated by & and its absolute convex hull is the same (without loss of generality), we additionally assume the following:
Assumption A.2. is absolutely convex.
Remark A.3. We note that the assumptions A.1 and A.2 are needed only to generulize our theorelical results an IPM-rreularized UOT formulation (Formulation 13). These assumptions are satisfied whenever the IPM employed for regularization is the MMD (Formulation with kernel that is continuous and universal (i.e., c-universal).
A.2 Classical Examples of IPMs
Maximum Mcan Discrepancy (MMD): Lot be a charactoristic kernel (Sriperumbudur et al. 2011) over the domain X, let [| denote the norm of f in the canonical reproducing kernel Hilbert space (RKHS), Hx, corresponding k. MMD, is the IPM associated with the generating set: G1 (f Hal [[flle < 1).
/ dso J dio
MMD 40)
max FE9x
Kantorovich metric (K.)): Kantorovich metric also belongs to the family integral probability
max If(x))f(w) metrics associated with the generating set W. = { X =EX*VEX <(m.g)
$ $ where
is metric over X. The Kantorovich-Penchee duality result shows that the 1-Wasserstein metric is the same as the Kantorovich metric when restricted to probability measures.
Dudley: This is the IPM associated with the generating set: Da {f: X > IR || + l|flla I 1), where d is ground metric X X. The
so-called Flat metric is related to the Dudley metric. It's generating set is: Fa x ll I 1, 1}.
Kolmogorov: Let X' = R". Then, the Kolmogorov metric is the IPM associated with the generating set: K. = (1(00)))) R*).
20

---

Published in Transactions on Machine Learning Research (01/2024)
Total Variation (TV): This is the IPM associated with the generating set: T ( + R |/ I 1)), where || = max
sex [/(z)] Total Variation metric over measures 80, t0 E
(x) defined as:
s(y) ((y) t(y) s(y) otherwise
if sf) (yy
[V,,,t) = Jy d|s t((v), where Is (s))(y))=
B Proofs and Additional Theory Results
As mentioned in the main paper and Remark most of our proofs hold even with a general IPM-rogularized COT formulation (13) under mild assumptions. We restate such results and give a general proof that holds for IPM-regularized UOT (Formulation 13) of which MMDD-regularized UOT (For- mulation 6 is special case.
The proposed IP IPM-rgullrized UOT formulation is presented as follows.
/ dı + A so) + A2
Ug (g to)
min =x
(13)
where "g is defined equation 12
We now present the theoretical results and proofs with IPM-regularized UOT (Formulation 13 of which MMMD-regularized UOT (Formulation 6 is special case. To the best of our knowledge, such analysis for IPM-regularized UOT has not been done before.
B.1 Proof of Theorem
c E C(X' x X) and X is Theorem 4.1. (Duality) Whenever G satisfies Assumptions A.1 and A.2
compact, we have that:
Ug.addid (80, to)
max (E
Sx f dsp + fx 9 dta,
s.t.
f(x) + g(v) S c(z, v xxyex.
(14)
Proof. We begin by re-writing the RHS of 13 using the definition of IPMs given
12
[
/ 9 drz
Ug (80,
min sER (X'x2)
dx + A
max TE9
f dso
f dm
+ A2
max 9E9
dto -
xxx
x
X
x
x
[
/ f dm
min (xX-(X(X)) XxX
dr
+
max FEG
dso
+ +
max DEG
9 dto
9 dnz
x
I
J g doo
min =ER (Xxx)
dr +
max SEG()1)
f dso
fx4 di
max (EG((2)
g dto
XxX
x
x
/ XxX
dn f f dri
f g drT
max (EG(S)))EE(Aa))
f dsu
g dto
min *x
fx
x
f c- dr
max SEG(A)))ES(AA)
[ dss
g dlo
+
min =x
x
xxx
+ { 0 if f(x)) +Q()) S vl v.
Here, F(x,y) = f(x), g(z. v) = g(v). The min-max interchange in the third equation is due to Sion's minimax theorem: (i) since R(X) is topological dual of C(X) whenever X compact, the objective is (ES(A)))
max
fx
f dso
g dto
-c0
otherwise.
/ |so
max SEG(A)))EE((AA)
g dlo.
x
s.t. f(x) + glv)) S ((I
(15)
21

---

Published in Transactions on Machine Learning Research (01/2024)
bilinear (inner-product in this duality), whenever c, are continuous. This is true from Assumption A.1 and E x()))). (ii) one of the leasibility sets involves 9, which is convex compact by Assumptions A.1 A.2 The other feasibility set is convex (the closed conic set of non-negative measures).
Remark B.1. Whenever the kernel, k, employed is continuous, the generating set of the correspording MMD satisfies assumptions A.2 and G& C C(X). Hence, the above proof also works in our case of MMM-regularized UOT (i.e., to prove Theorem 4.1 in the main paper).
B.2 Proof of Corollary 4.2
We first derive an equivalent re-formulation of 13 which will be used in proof. Lemma B1.
Ug (o) = min (2)
(s,W + (0) + (t,
(16)
Wi ( ( if [s [t), otherwise.
where Wi (s,t)
with W as the 1- Wassersteii metric.
00
Proof.
min (x)
s1) + (2)
dit + Arrglss so) + (t, s.t. T / cdT
8
min +.))+(()) |6|=||| |s|
min (xX](x(x))
T2
[s]
8
/ cd + \r so) + (1)2,,)
min 1/5
min (2xx)
min 120
min aER (1xx))
(ndd Xino(1F 80) (2
/ / + A17g(TT 80) A2 (ER=(XxX)
min
W + 8 otherwise.
[s 109=11]] |t]
The first equality holds from the definition of Wi: Wi(s, =
Eliminating
normalized versions and t using the equality constraints and introducing n denote their common mass gives the second equality. The last equality comes after changing the variable of optimization to æ e RR (X' x X) nt. Recall that R R*()) denotes the set of all non-negative Radon measures defined over X: while the set of all probability measures is denoted by R (2)).
Corollary 4.2 in the main paper is restated below with the IPM-regularized UOT formulation by its proof.
13 followed
Corollary 4.2. (Metricity) In addition to assumptions in Theorem 4.11 whenever c is metric, Ug belongs to the family of integral probability metrics (IPMs). Also, the generating set of this IPM is the intersection of the generuling set of the Kanlorovich metric and the generating set of the IPM used for regularization. Finally, Ug is valid norm-induced metric over measures whenever the IPM used for regularization is norm-induced (e.g. MMD with a characteristic kernel). Thus, U lifts the ground metric c that over measures.
Proof. The constraints in dual, (7), are equivalent to: g(y) S min c(z, v) f(x) Yy X. The RHS is zEX nothing but the c-conjugate (o-transform) f. From Proposition 6.1 (Peyré Cuturi 2019 whenever f(y) if f f We. Here, W. is the generating set of the a metric we have: min c(,, y) f(x) -o0 otherwise. =EX
Kantorovich metric lifting Thus the consts equivalent to: glv) S -(() V W..
22

---

Published in Transactions on Machine Learning Research (01/2024)
Now, since the dual, seeks to maximize the objective with respect to 9, and monotonically increases with values of 9i at optimality, we have that s(y) -f(v) A X. Note that this equality is possible to achieve as both o--f G (these sets are absolutely convex). Eliminating 9, one obtains:
Ug (80,
f dso Sx f dto. JEG(A)MW
Comparing this and the definition of IPMs 12 we have that Ug C,A.A belongs to the family of IPMs. Since any IPM is a pscudo-metric (induced by semi-norm) over mcasures (Muller 1997], the only condition left
to be proved is positive definiteness with (5,,.2.1800,40) 40). Following Lemmn B1 we have that for optimal s", tt in [16), Ug.c,x.x50, to) =0 = (t) Wi(s",t") -+,+,====0 = 0, (ii) "g(s", sp) 0, (mi) Yo (t", = as each term in the RHS non-negative. When the IPM used for regularization is a norm-induced metric (e.g. the MMD metric or the Dudley metric), the conditions (t), (ti). (iii) + s* = so = to, which proves the positive definiteness. Hence, proved that Ug.., is norm-induced metric measures whenever the IPM used for regularization is metric.
Remark B.2. Recall that MMD is valid norm-induced IPM metric whenever the kernel employed is char- acteristic. Hence, our proof above also shows the metricity of the MMD-regularize UOT (as per covollary 4.2 in the main paper).
Remark B.3. If G is the unit uniform-norm ball (corresponding TV), our result specializes to that in (Piccoli & Rossi, 2016), which proves that Ug,., coincides with the so-called Flat metric (or the bounded Lipschitz distance..
Remark B.4. If the regularizer is the Kantorovich metrid) i.e., G = Wo, and d = = then U,X. coincides with the Kantorovich metric. In other words, the Kantorovich regooooricc-rgglariieede OT is the same as the Kantorowich metrie. Hence providing an OT interpretation for the Kuntorovich metric that is valid for potentially un-normalized measures in R+ (X).
B.3 Proof of Corollary 4.3
Proof. As discussed in Theorem 4.1 and Corollary 4.2 the MMMD-regularized UOT (Formulation 6 is an IPM with the generating set an intersection of the generating sets of the MMD and the Kantorovich- Wasserstein metrics. We now present special cases when MMDDregularized UOT (Formulation recovers back the Kantorovich- Recovering Kantorovich. Recall that G& (A) = (Ag | E 9x From the definition of (x),, / G. f H4. || & A. Hence, as À (A) = Hg. Using this in the duality result of Theorem 4.1 we have the following.
metric and the MMD metric.
/ fdso - fdto=
/ fdso - -f fdtt
lim UR.., to)= lim max A
max SERunW.
x000 (e
/ [ddso -1 fdtt
max SEC(X)O.
max feW.
fdso - fdlto
Equality (1) holds bocause M4 is dense in the set of continuous functions, C(X). For equality (2), we use that W. consists only 1-Lipschitz continuous functions. Thus, Vsg, to E R (X), lim Ui,exx, 40) (s to).
Recovering MMD. We next show that when 0 < A2 A 5 1 and the cost metric is such that
c(x, y) 2 =(x k(y, y) 2k(x, v) [[ő(x) - d())ll& Vz, y (Dominating cost assumption discussed
in B.4 then Vsa. to E R+((2), Ms. 4o) MMD 4o)
3he ground metric Mg..A.A. must be the same as that defining Kantorovich regularizer.
23

---

Published in Transactions on Machine Learning Research (01/2024)
Let f (k(A)) f= Ag where g E Hk. |g|| 1. This also implics that Ag He A E (0.1).
1f(x) (ww)) = (Ag, ((x) <(v)) (RKHS property)
- (g, x(x) - +(v)) (=0kA=ı))
< lg)lall((x) <(y)] (Cauchy Schwarz)
< [[(xz) - <(v)lla (:: ||9|| S 1)
s c(2, y) (Dominating cost assumption, discussed in B.4
f We
Therefore, (A) We and hence, (2) (A)... This relation, together with the metricity result shown in Corollary 1.2 implies that Uk,c,A,A tu) = (s, to). In B.4] show that the Euclidean distance satisfies the dominating cost assumption when the kernel employed is the Gaussian kernel and the inputs lie on a unit-norm ball.
B.4 Dominating Cost Assumption with Euclidean cost and Gaussian Kernel
We present sufficient condition for the Dominating cost assumption (used in Corollary 4.3 to be satisfied while using Euclidean cost and a Gaussian kernel based MMD. We consider the characteristic RBF kernel, k(x,y) y) = exp (->lx y)l²), and show that for the hyper-parameteer, O < 8 S 0.5, the Fuclidean cost. is greater than the Kernel cost when the inputs are normalized, i.e, ||>||| = || = 1.
o y| 2 k(x,x) + k(y.8) 2k(x,8) - |0|| + ||y12 2 (α. v) 2 2-
- 2k(r, v)
(17)
(z, v) v exp ( 2s((1 - ()) ( (Assuming normalized inputs)
From Cauchy Schwarz inequality, ||ا|||||||| < (z, y) < v|| With the assumption of normalized inputs, we have that -1 S (z, v) < 1. We consider two cases based on this.
Case 1: (z, v) [-1,0] In this case, condition 17 is satisfied Vs because k(x,y) 0 Vzz y with a Gaussian kernel.
Case 2: (z. (0, 1] In this case, our problem in condition 117 is to find 0 such that In (z. v) S -2s(1 (I. We further consider two sub-cases and derive the required condition as follows:
Case 2A: (z, v) E (0, 1
We re-parameterize (x, () = e for n 2 1. With this, we need to find s 2 0 such n 2 2s(1 - -=)). This is satisfied when O < 8 & 0.5 because e-n 2 1 - n.
that -n < 2s(1 - e n)
Case 2B: (x, v) (2, 00) We re-parameterize (x, v) = e for n 1. With this, we need to find s 2 0
such that
2 2s. We consider the function f(n) = n e = for n 2 1. We now show that f is
1-0
increasing function by showing that the gradient
df = (1 1 e is always non-negative. dn
dʃ
2 0
dn
e#
1 +
In
1 + 20 0
(In
(n + +1) - In (m)) 20
24

---

Published in Transactions on Machine Learning Research (01/2024)
Applying the Mean Value Theorem (n)) In n. get
n) 1
In (n 1) In
(n +
where = n
In + S n
2
df in 1+ 20 0 dn
The above shows that f an increasing function n. We note that lim f(n) = 1, hence, T(a) 2 which implies that condition (17) is satisfied by taking 0 < S 0.5.
1-1--
B.5 Proof of Corollary 4.4
Corollary 4.4 in the main paper is restated below with the IPM-regularized UOT formulation by its proof.
13
followed
Corollary 4.4. UG,,AA((s5) S min (Arr(ss (c,,
Proof. Theorem 4.1 shows that CgA.A. is an IPM whose generating set the intersoction of the generating sets of Kantorovich and the scaled version of the IPM used for rogularization. Thus, from the definition of max, we have that Ug,,,x,s,, I) S (n 1) and Ug...AA(s, s K.,(s,) (). This implies that Ug,,,X.A858)) min (s,
Ko,s,, (s, f)). As a special case, (s, I min (AMMD& (s)
B.6 Proof of Corollary 4.5
Corollary 4.5 in the main paper is restated below with the IPM-regularized UOT formulation 13 followed by its proof.
Corollary 4.5. (Weak Metrization) egA. metrizes the weak converyence of normalized measures.
Proof. For convenience of notation, we denote Ug,,,,,1 by U. From Corollary 4.4 the main paper,
< U(3,,, B) S Ke (3n, 3)
From Sandwich theorem, limg,-9U(3.,. 3) as limg, 3 K. (Bn, B)) - 0 by Theorem 6.9 in Villani
2009
B.7 Proof of Corollary 4.6
Corollary 4.6 in the main paper is restated below with the IPM-regularized UOT formulation by its proof.
13
followed
Corollary 4.6. (Sample Complexity) Let us denote Ug.e.x.>, defined in by U. Let 8m, fm denote the empirical estimates of 8o, to R- (X) respectively with samples. Then, Ts fm)) U(s0; 1o) at rate (apart from constants) same that of (8 so) + 0.
Proof. We use metricity of U proved in Corrolary
4.2 From triangle inequality of the metric ú and Corol-
lary 4.4 in the main paper, we have that
0 0 ((3m, (m)) (30. lo)l < ((m, so) + (1o, 1m) < X (7 80) + 7a (fm, lo))-
Hence, by Sandwich theorem, U( (m.) + Ū( (o) at rate at which 7o (5mm so) - O and 70(6m
to) O. If
the IPM used for regularization is MMD with a normalized kernel, then MMD3 (so. sm) S v m with probability at least 1 1-6 Smola et al. 2007).
2 log(1//
m
From the union bound, with probability at least 1 5, 127 (8m- (m) - (80, 21 (V +
(log(288)
25

---

Published in Transactions on Machine Learning Research (01/2024)
B.8 Proof of Theorem 4.8
We first restate the standard Moreau-RRockafellar theorem, which we refer to in this discussion.
Theorem B2. Let X be a real Banach space and f.g X + R U foo} be closed convex functions such that dom(f)ndom(g) is not empty, then: (f+g)* = min æi #
(x1) Vy E X". Here, f is the Fenchel
conjugate of f, and X is the topological dual space of X.
Theorem 4.8 in the main paper is restated below with the IPM-regularized UOT formulation 13 by its proof.
followed
Theorem 4.8. In addition to the assumptions in Theorem 4.1 if e is a valid metric, then
dg (s0, (o) = min K.(s, () + A 80) + A2 s.ER(()
(18)
Proof. Firstly, the result in the theorem is not straightforward and is not a consequence of Kantorovich- Rubinstein duality. This is because the regularization terms in our original formulation [3 16 enforce closeness to the marginals of a transport plan and hence necessarily must be of the same mass and must belong to R Rt(( (X). Whereas in the RHS of 18 the regularization terms enforce closeness to marginals that belong R(x) and more importantly, they could be of different masses.
We begin the proof by considering indicator functions Fe and Fg defined over C(X) C(X) as: 0 if f(x) + g(v) S c(r,w) V x, 1 E X, .FAA,,
Fe(S.g)= = {
0 if ( E G(22),
0o
otherwise.
oo
otherwise
Recall that the topological dual of C(X) is the set of regular Radon mcasures R(X) and the duality product s) =ʃʃ ds v ʃ C(X), E R(X). Now, from the definition of Fenchel conjugate in the (direct sum) ds + f dt, dt,s.t.f(z) ( I ((ry) V z.y E X, max J f where 8, E R(X). Under the assumptions that X is compact and e is continuous metric, Proposition 6.1 in Peyré & Cuturi 2019) shows that (s, max f fds - J fdt = (s, JEW.
space C(X' @ C(X), we have: ts, fEC(x))
On the other hand, P2 = max I ds - max I g dt Angg(s,)) + (e 0). Now, (E9(Ai)
min gE9(A2) we have that the RHS of 18 (o)
(s, + (1,, This be-
cause ( 0, (Sp. (S..s)).. s). Now, observe that the indicator functions Fg Fe are closed, con- closed, convex sets. Indeed, 9 is closed, convex set by As- sumptions A.1 A.2 Also, it is simple to verify that the set {(f.g) f(x) + g(v) S Se(z, v) V Vz,y X3} is closed and convex. Hence by applying the Korcau-Rockafella. formula (Thcorem B2], we have that the RHS of 18 is equal to (Fe + Fe.drig)* (sp, lu)- But from the definition of conjugate, we have that (Fa + )2))
vex functions because their domains
(Sa, fo) = max Sx d8o + fx g dtg - Fe(S.g) (f.A Finally, from the fEC(()),EC((x)
definition of the indicator functions Fe, Fa.A1.A this same as the final RHS in 15 Hence Proved. J
Remark B.5. Whenever the kernel, k, employed is continuous, the generating set of the corresponding
MMD satisfies assumptions A.1 A.2 MMD-UOT.
and Gr C. C(X). Hence, the above proof also works in case of
B.9 Proof of Theorem 4.10 Consistency of the Proposed Estimator
Proof. From triangle inequality,
Vmm fm) ((s0 S K4m fm)) - Um
m (s0, to) + 12 (.. to) - ((0. to)l: (5.
(19)
where Um (sn, to) is same as (s0, 40) except that it employs the restricted feasibility set, F(8m (m), for the transport plan: set of all joints supported using the samples in 8m, Em alone i.e.,
F(8ms E. and (mn))))))))
Oij 2 2 1, ....................... m2 Here, 82 is the Dirac measure at We
begin by bounding the first term in RHS of 19
26

---

Published in Transactions on Machine Learning Research (01/2024)
We denote the (common) objective in tmm (.,)),,((()))) Then,
function of the transport plan, n, by h(T,,;)).
Uml fim) (m0
min xEF = (mm.tm)
11( 8m fma)
min (fm(m,
N(T, 80, t0)
S Nu(TT 8ms (mm) 1 11x ,0, where TH arg min h(T,, 80, to) xEF(lm.fm)
= A1 (M 3m) - MMD (xf" + A2 (MMD (TT)" (m) - (T
S ArMMDA(SO. S.) + HcMMDA(to. fm) (:: MMDs, satisfies triangle inequality)
Similarly, one can show that Um (s00 to) (8m. 2m) S M 8m) + A2MMMD fm) Now,
Muan-
det et al. 2017, Theorem 3.4) shows that, with probability at least 1- MMDx( 6m) m
log(1/5)
where & is a normalized kernel. Hence, the first term in inequality (9 is upper-boundee by (A1 + X2)
Vm V
2 2/5 with probability at least 1-
m
We next look at the second term in incquality 19 [2 (s0, to) - (s0 to)l- Let nth be the optimal transport plan in definition of (su. to). Let x' be the optimal transport plan in the definition of ((s. to). Consider another transport plan: Am E F(8m,mm) fm) such that *''(x; plang)) where 7T(;,, Mj) x (mt.y.) for i.j E m1.
so(1r))o(V5)
A4m (su, 40) Ü(s, 40)l (80, 1o) ((0, 4o)
h(=",, s0, to) - h(T=" Su. to)
S h(*m -8. to) h(T= so,tto to) (:. nth is optimal.)
< c dam - / di + AIll (mi))ll + dellun(ATT)) pu (T)))]ll
(:: Triangle inequality)
To upper bound these terms, we utilize the fact that the RKHS, H4: corresponding a c-universal kernel, k is dense in C(X) wrt. the supnorm (Sriperumbudur et al. 2011 and like-wise the direct-produet space, Ha @ Ha, is dense in C(X' X) (Gretton| 2015). Given any I E C(X) x C(X), and arbitrarily small >> 0, we denote by f-e the functions in H4 & H4 that satisfy the condition:
- e/2 f ffsfsfftee
Such an fe Hy He will exist because: i) f e/4 E C(X' C(X) and ii) Ha He C(X) C(X) is dense. So there must. exist some Se He @ He such that [s v) e/4 (r 5 x, E X b f(z v) < fe(x, < f(r +e/2 2, X. Analogously, f-e exists. In other words, fe. H4 & H₂ are arbitrarily close upper-bound (majorant), lower- (minorant) of f C(X) x C(2)).
We now upper-bound the first of the set of terms (denote so())to(v) by E(x, v) and (T is the corre- sponding empirical mcasure):
f d - f dr < d
-f
c-e dr
(Ce, (a
(c-c)(M((TT))
-Ce, H4: (=mh)
(ce,HH((T))) (ce, px (c-e:
(A Mx(TT)) + (ce - c-es ((e")) 5 (CE, (a") (x()))) + -.
(:: ||c.- - -c-elloo < and define a, as the mass of measure s) $ (cc|lxll/4<( - (TT co
27

---

Published in Transactions on Machine Learning Research (01/2024)
One can obtain the tightest upper bound by choosing ce arg ue ||ell Accordingly, we replace ||c|l by g(t) in the theorem statement Further, we have:
s.t. c c v c c/2.
[|)=(**)) - /x<(==)] = /1/ dz (x) @ (x(y)d"""xa - on (x) @ (x()d) v)
/ (x(()) @ (x())) x' (x,
on 2 (x x (cr u)), (o(e') @ on (F x ((d),))) (x) on (5T ((c. I da((") s onll (sm (3' (((d:))
/ / (oL()) ( / (6L(1), oh / f k(z, x'),'')( vf m1 v)m(e' s (z.))) E(z,8)) d (ém v/) E(a'.
E(x', ./))
(2, v)n('' y'dd (x) ((..v) a ((""('','))) - E(xl,1"))
- Mn v)n(e' s'dd (z,));)) E(., v) a - 2(1).2))
Now, observe that k X x X x X' x X defined by le ((z, 4) (z/, = k(x, x') k'y, yf (mx kernel. This because h = kakiohic where ka ((x, u): 2/)) = k(r,, a") a kernel, hy ((x, y), w') = k(y 3²) is a kernel, and & ((z, y), (x', (z),8/)) 17(2, y)m(z', ./) is kernel (the unit-rank kernel), and product of kernels is indeed kernel. Let (xx, y) be the feature map corresponding to h. Then, the final RHS in the above set of equations is:
y)q(2', ((o'',, v/) ) is a valid
=
/ (1ba (a))))) v). (0d' w') (z,)) ((t..)) ((r.w)) a (fm v) s(4',5))
( (x w)d (5m (x,, v) ((..w))) / (a)a (x',v) E(d', v/))
Hence, we have that: ( (m((T)))l = | -
Again, using Muandet et al. 2017
Theorem 3.4), with probability at least - ō, | | -<a)))
i S Ca + /2C² log(1/8) where C3 =
((x, v), (/),)/)). Note that O6 as X is compact and so, to are assumed to be positive measures and ls is normalized.
max.
Now the MeD-ragulrrizer terms can be bounded using a similar strategy. Recall that, (x;) (a),)))) E Holo(z)) so we have the following.
/. (c(r))d((()) - f (x(r))di(()))
1 (+(") - 1 ( =
// (e(r)) (x) (000-(((x)) x(())) k
= / (x))) - +i(x)) / (x(=')d - -aj(a'))
= / (x()))(a))) w)d (' v) - v(c.v)) / "ola')((a' s/d (ém(Z', 5(2' (((a
/ (on(x), (h (2)))) v)rla' s' dndd (em(x..) - ((((,.)))((x))))) d ((m(sr E(x',
E(x', .s))) / Ma,,z'm(a,, w)n(d' an)d (x,,)) (("((,,)))(( -(xrs))) d ((TT( (dm
4This leads to slightly weaker bound, but we prefer it for ease of presentation
28

---

Published in Transactions on Machine Learning Research (01/2024)
Now, observe that 1: X 2' X X' defined by k ((x, (,,, (z', k(xx)'))(z (on(an)n()) 3/) is valid ker-
nel. This is because k = kn k2, where ky ((x, v): (a' v/) k(ar,=') is kernel and k2 ((x, u). (x',
nf, v)m(m'/,/))) v') is kernel (the unit-rank kernel), and product of kernels is indeed a kernel. Hence,
we have that: ( ( = |g((TT) - (x(
Similarly, we have: I1 - Ms (T2)) =
( 1z(5)
Again, using [Muandet et al. 2017 Theorem 3.4). with probability at least - 8,
(ax(5")) -x((E)
max & ((z, u), (z',8')) Note that C S Cs + /20, log(1/6) where Cß ".y.a' y''E
X' is compact, 80, to assumed to be positive measures, and k is normalized. From the
union bound, we have:
(m fm) ((s, to) I (A1 + A2)
1 vm
+
log (5/6)
Cs. + V2C, log(5/6)
+
g(c) C m + 22Ck log (5/5)
+ c0x; with probability at least - 8. In other words, w.h.p. we have:
Üm (8ms bm) - (so,
S 0
Arttd vma
+ g(t) <0.))) for any > 0. Hence proved.
B.9.1 Bounding g(e)
Let the target function be approximated be he E C(X) (x),, which is the set of square-integrable functions (wrt. some measure). Since X' compact, being c-universal, it is also C²-universal.
Consider the inclusion map N4 (X)), defined by g g. Let's denote the adjoint of & by t.. Consider the regularized least square approximation of h² defined by h (c* + 1"h E Ha, where > 0. Now, using standard results, we have:
|1ch h² 12llc> (((044))76
= -[((((((+)))-- (e + t)( + + lee
The last incquality is true bocause the operator L" is PD and t 0. Thus, if = f= |l,h; lo I Ilcha lellca Clearly,
I1( -- then
g(c) < |14|ll4
(he he.
V /((e+
(10(
V((e++ (0l0++))-e0(4440 +) 2/1 -/,, h1))
= V((rr)) (1 hr, (1 00) (8)+
= -1( (e + +t - 1|
xá X+i This maximized when À A=6 Hence, f(A) 1 1. Thus, 2/8
Now, consider the spectral function f(A)
g(E) ||A" [22 1 1 13
Therefore, decays as m473 then, g(c) O m//33
2/8
29

---

Published in Transactions on Machine Learning Research (01/2024)
B.10
Solving Problem (9) using Mirror Descent
Problem 9 is instance of a convex prograi and can be solved using Mirror Descent (Ben-Tal & Ne-
mirovski 2021 presented in Algorithm
Algorithm Mirror Descent for solving Problem
Require: Initial a1 2 0, max itcrations N. f(a) TT: (aC12) + A ||a d all... + A2 Ia
G2
for - 1 to N do
if [V * 0 then
$ = 1/ V/(a)ll
else
return ag.
end if
0441 04 ① - (.V((())
end for
return a;+1.
B.11 Equivalence between Problems 9 and (10)
We comment on the equivalence between Problems forms:
and 10
based on the equivalence of their Ivanov
Ivanov form for Problem
99 is
min
Tr (aC s.t.
al 01
T 1
02
1
S r2.
a20ERm3 Xmn
m
Gn1
m2
Gz
where r1,72 O.
Similarly, the Ivanov form for Problem
10 is
2
2
min a20ERma
Tr (aCT2) s.t. al
1
S f1
T 1
02
r2.
m1
Gn
m22
Gz
where F1,72 0.
As we can see, the Ivanov forms are the same with 71 = =ri, F2 = =r2, the solutions obtained for Problems 9 and 10 are the same.
B.12 Proof of Lemma
Proof. Let f(a) denote the objective of Problem 10 target samples, respectively and m1, m2 as the number of source and target samples respectively.
G11, G22 are the Gram matrices over the source and
Vf(a) C12 + A1G1 (olm²
1ms
I + delm, (Im, 1 ma
Gz2
mi
We now derive the Lipschitz constant of this gradient.
Vf(a) vf()) = 2 (a- m I 1m, 1 de (a B) Gzz)
vec ((Vf(a) ()())) 2 2(Aivec (G1 (a B) 1 1m 1 T + dzvec ((1m 1,,, (a 3) G2a) G2a) )))
= (A11mm & Gn + A2 8 1m 1mT)) vec(a - 3)
where denotes Kronecker product.
30

---

Published in Transactions on Machine Learning Research (01/2024)
(V)(a))(a) -f(())] l|vec ([Vf())VV(33)) 1|8
S 2 I 8 Gi + X2G 8 mm 1 ||]]vec[[r B) (Cauchy Schwarz).
This implics the Lipschitz smoothness constant
L 1 mng @ G.1 + 22 mm 1 ||
= (2/()m2)))IIn +
+ 221 ( (1mg I & G11, G22 & Lmy 1
= (Glll) + (A2m1 G22l + 22 Gnil;; (I Gzzlmz).
For the last equality, we use the following properties for Kronecker products- Mixed product property: (A @ B) = & BT, (A @ () & D) = (AC) @ (BD) and Spectrum property: Tr ((AC) & (BD)) = Tr (AC) Tr (BD).
B.13
Solving Problem
10 using Accelerated Projected Gradient Descent
In Algorithm
we present the accelerated projected gradient descent (APGD) algorithm that to solve
Problem [10
as discussed in Section 4.2
The projection operation involved is Projectzo(x) max(x, 0).
B.14 More on the Barycenter problem
B.14.1
Proof of Lemma
Proof. Recall that we estimate the barycenter with the restriction that the transport plan स² corresponding (A,, s) is supported on D; x D Let 3 Thm denote the probabilities parameterizing the barycenter, s. With Um as defined in Equation (9.. the MMD-UOT barycenter formulation, (m, 8,)) = min E pillm (8i, (B)), becomes 820
n EalTee
Σ pi (TTTT(CCT))) (a,C)) + A1Ila;1
0i Illa + A21 1 - BlGG m,
min .20
(20)
Following our discussion in Sections 4.2 and B.11 we present an equivalent barycenter formulation with squared-MMD regularization. This not only makes the objective smooth, allowing us to exploit accelerated solvers, but also simplifies the problem, as we discuss next.
By ($1, ... 3w)
min
Σ p Tr (a(c)) + A1lla - z 111.. + A2laa 312
(21)
920
The above problem is a least squares problem in terms 8 with a non-negativity constraint. Equating the gradient wrt. B as 0, we get G(B ΣØ= p 1) = As the Gram matrices of universal kernels are full-rank (Song. 2008 Corollary 32), this implies 3 = Σラ= p3aT. 1, which also satisfies the non-negativity constraint. gives us the MMD-UOT barycenter formulation:
Substituting ΣØ= p J
1 Ba (81, ... 8n) min
Σ t=1
Tr (r(ctt)) A A1/1 1 112.. A21la 1 m;
Σ j=1
pga3 112
(22)
320
31

---

Published in Transactions on Machine Learning Research (01/2024)
B.14.2 Solving the Barycenter Formulation
The objective of 22 as a function of Ois, has the following smoothness constant (derivation analogous to Lemma 4.11 in the main paper).
L, 2p= (Amm)> + (nmm;)² ||G|] + 21 G
where Tp = A2(1-114 pi). We jointly optimize for o using accelerated projected gradient descent with step-size 1/L.-
B.14.3 Consistency of the Barycenter estimator
Similar Theorem 4.10 we show the consistency of the proposed sample-based barycenter estimator. Let
8, be the empirical measure supported over samples from 8;. From the proof of Lemma 4.12 and 22 that,
recall
B (8), ... .sa)= min Σ pi Tr (a, + A A1lla111 3:lle,. + >21la 1
Σ pja3 T 118)
Now let us denote the true Barycenter with squared-MMD regularization by B(s). ,n) = min LisiP (x)
where U(s = min
Sc dm' + A1M (=), st) X2MMD s). Let
E'R (2)
x ,,...,"h - 8 be the optimal solutions corresponding to B(s,,".. s,)) It is casy to see that s"
Ej=1 PyTT"" (for e.g. refer (Cohen et al. 2020 Sec C)). After climinating 8, we have: " = min EX)) E Ps ( M,D (Tf. 8;) A2MMD Ej ()T)))
Theorem B3. Let 1/ 2x) = (3) where is the misture densily = Ei= # Under mild assumptions that the functions, n'cc He @ Hr, we have that w.h.p., the estimation error, B ($1, 3m) B(. sn)l O( max (Ibr'Ilalee())) /m).
vE]1.m)
Proof. From triangle inequality,
1B', ($1, ... B ([(s1,.... ... ,x)1 & [B,, ($1. ... ,8,) BB) (81, ......... sn)]+ ......... (... -B,))((s... " ,8)),, (23)
where Bin (81, ,,) is the same as B( ,n) except that it employs restricted feasibility sets, F(($1. ... ,8)) for corresponding a as the set of all joints supported at the samples in $1. Sn alone. Let D. = -xi,, ".. 3mm) and the union of all samples, UDR 1-1 = (=1... ... }mm)
Fi 3n)
(ETh Σ/ ((f((su))))) an 2 -1A-= I...., m:j = I....,,,| Here, 8, the Dirac measure at We begin by bounding the first term.
We denote the (common) objective in (.), B(-) as function of the transport plans, (T ""),,, by h(TT! ..
Bi (51, ... 3.)) By, (s,, ,8))
min h(T TT, ,.,)
min FF (a))
h( ... x² $1, ,n)
FFE =* $1, B..) h( _ $1. 8n)
.n)
h(
where ~ arg
(1 ..)) h(T² T" 81, .n) for i [1,
dap (MMD) ( 8,) - ( 8:)
#= Pidi (MMD (T 3.) - MMD st) (MM MMD (s)*
2 >n - Ps (MMDx (..) MMD s:) M Ei P&MMD (&,, s;) (As MMD satisfies Triangle Inequality),
2A, M max MMD 8;)
sE[1.a]
32

---

Published in Transactions on Machine Learning Research (01/2024)
where for inequality (1) use that compact.
(s,t(sst) M < 00 the generating set of MMD is (.T)
As with probability at least 6, MMD 8:)
2 llg(1/6) Smola et al. 2007:: with union
vm
m
bound, we get that the first term in incquality (23) is upper-boundee by M probability at least d. 5.
vm
I
+
2 log 2n/6
with
We next look at the second term in incquality (3) : B ($1. ... ,,,) (ss s8)).. Let (T' x²) be the solutions of B (s1.... sn). Let (x ,... )") be the solutions of B(si, 8n). Recall that s' denotes the mixture density s' = = st. Let denote the empirical distribution of s' by 3' (i.e., uniform samples from - D:). Consider the transport plans: *tm Fs 88) such that atm((1,) ('(m5,) where m
33

---

Published in Transactions on Machine Learning Research (01/2024)
2z), = ( for I E [1, E mn]. s(()))=(3)))
18m((1) (.1, ,8,) B(si,.... sn)] (81
,8,) 1 BB(11, ... )n.)
11(ā 1m
7 mm 81, 8n) 11( n 81, :.. kn)
S h(a
Amm,, 81, sn) h(a T .1, 8,)
,
Σ Pi
(mua(ttmm) um( ca) + 2A1 M (+(m))) - uu (s:)lla
1=1
[((x((+1")) - (x((s))l))) (x)))
EerTst 11)
(m((t(")) - pm E n
2A2M
Σ
fm Pg 22
(Upper-bounding the sum of two MMD terms by 2M)
< Pi (uo( [x((T")), c;) + 2A, || - (m(((())) 11
1=1
PogT 11)
(mn(tt37)- MA Σ
( - up Σ n
2A2M
25 jm
(Using triangle inequality)
Σ Mi
(ma(attm) ma(==))))))))) ),()) + 2A₁ M Ml)m4((fTm)) (=((="")ll
[Iua(E?T)) (u(TTz))l 1| Σ Pollux pa ( ( Ill
2AaM
(Triangle Inequality and linearity of the kernel mean embedding)
n
S Σ Pi
- ( [ ||c|||
2 (lan(("") - pn (m(""t I
[|ux(((m)) me((T"") Σ Pollux - (A(T)")ll)) 211m) )
M
(Cauchy Schwarz)
< max sE(1.x)
Il (a(m) (T(*")) l|calla + 2A MA (llx+ +
(lea((85m)) un ( + max [|am(()"m) un(=5"))l
2M
JE[l.n]
34

---

Published in Transactions on Machine Learning Research (01/2024)
We now repeat the steps similar to B.9 (for bounding the second term in the proof of Theorem 4.10 and get the following.
(* tek (T((T"")) )|k
[ f d# - S f dr
max SEH Sll
max
f d f dm'
/EHu.l fllk
max
m 1=1
mn 2-=1 f( 21))
m
[ 2)
(x(x)s(22) dx dz (.))) (.)
SEHu.l $1
f(x, 2) (X.)
max
Ex 8, 7,
SEHH-l S1
max
[f(X, 2)ny' ZX
/EH2il fllk S1
(S n' ((),((XX00())) ( (()))
max 1x, 7,,8
FEHu-l f1
max
(f 2 1/, 2-
1
SEH2. S1
max
@ -
JEH2.l $1
(:: Cauchy Schwarz)
max
|I.fll/l ||x [Bx-A -
((Z)
SEHu.l S1
(:: properties of norm of tensor product)
max || |I |Bx-8-- ( @ (X)) 3 RZ (z( @ (()) lle
fEHu.l flleg1
||a||| Ex=e;=s; (x() x(X)] ||x E (o ((Z) ||a
Inr' ($a) (ea(s:) k2 I (5') Uka lla
(: o() @ 6() is the feature map corresponding to k2.)
Similarly, we have the following for the marginals.
If I
/ f d
(+(m) - -m((ii")ll==
dai imim
max 1lllal
/ f ddi
d
max fEH
x (2),, 2g
/ 11 f(x) (x,
max
(a
s((r)s(()) s1(2)s (z) dæ da
m2nsi (x1) (2g)
(x,
max
Ex-8,
1(X) 81 xi(X)((2)
fEHa.Ilfllass
max
Ex
[S(X)n (X, []]
FEHa-llfllass
Ex - nn 2~0 ( n/ o (xX o(x) (Z))
max
fEHa.llflaaS1
max
(f n2 Ex
(2)))
sEHe.llflaas1
max
& 1B
@ (Z)) lk
(:: Cauchy Schwarz)
max fEH
8 x(2) llx
(:: properties of norm of tensor product)
SE In | - (Ø(X) @ x(X) @ Ez~z-s' [o(2)) ll:
max
< |177///||xxxa-a0. (o(X) & x(X) | E [o(Z)) ||e
= 11r7 (8:) - ( / (5') - (=*ss))ls
(:. o(.) @ 6(-) is the feature map corresponding to k2.)
35

---

Published in Transactions on Machine Learning Research (01/2024)
Thus,
with probability at least
1
8,
1B (s1, ,8n) -B(s1: 8n)
max (e(l.m]
[In'lIl/le 2 |17''la + 2A2 (1lr II+ + max
11>/1la)] }) vm +
2og (2n+-2)/6
Ap-
plying union bound again for the inequality
in 23
get that with probability least
B (5 ,,)) -B(s1, "EL,,,,, (())l I vm
2 log (2n+-41/6
2A1 + Vm +
log (2n+-4)/6 )) where
5= max [Ir'lll + M + [E[L,]]
+ max JE[L,N]
||> (mll+t))
B.15 More on Formulation (10
Analogous to Formulation 10 in the main paper, we consider the following formulation where an IPM raised to the 4th power with 1 Z is used for regularization.
/ c AIn + (12, (o)
Ug (so, (o) = min =ER*(Xxx))
(24)
Formulation 10 in the main paper is special case of Formulation 24 when IPM is MMD and 2.
Following the proof Lemma B1 one can casily show that
d (so, to)
min S.ERI (X)
|s (s, + A1,8(s, 80) (t,
(25)
To simplify notations, we denote UG..A,A..2 by U in the following. It is easy to see that U satisfies the following properties by inheritance.
1. U20 0 each of the terms in the objective in Formulation 25 is greater than 0.
(o) -0 + 50 = fo, whenever the IPM used for regularization is a norm-induced metric. As Wi.ng non-negative terms, (s) for regularization is norm-induced metric, the above condition reduces to so to.
to) = 0 =
(g(5, so) 0, (t, to) 0. If IPM used
(so. to) Uto. so) as each term in Formulation
symmetric.
We now derive sample complexity with Formulation 24
Lemma B4. Let us denote Us.2dAdd2gg defined Formuulation 9 by U, where g > 1 Let 8m, im denote the empérical estimotes of so, to E RT (X) respectively with m samples. Then, U(Sm fm) + Uso. to) at rate same as that of pg(mm, + 0.
Proof-
(xxx) (0, / cd di + Xng 30) + Ang((T2)
(so,
min
We have,
U (8ms (m) (80,40)
min =ER+(Xxx)
M[T 8ms mm))
min (xxx)
1(TT, 80,4o)
< h(T", Sm, fm) h"(T", tp) where n" = arg min h( su,
FER (Xxx)
= A (T), sm) - - so) ) + 7 [T fm) - -
=> (10(7) 8m) (T]
2 (t Sm )m (1- 1 so))
(Z0(T), f.)) - (
Σ
7 ( ,g ( to)
36

---

Published in Transactions on Machine Learning Research (01/2024)
(mi, ng( 8m) 1 (o) )) +
70( 8m)
A
(ta fm)
7g (m f. 1 (m (n)
(:'79 satisfies triangle inequality)
S
To (80; 8m)
9 ( 8m)) 1- (m
l ( 80)
+
A (n0(to: fm) Σ
n (T2 Emm 1- ( f to)
= A (10(00- 8m) (70 8m) + ( 80))8-1
+ tto fm) (70( fm) r to)) )-1)
s X(2M)-1 (70(80, 8m) (g(, mm))
For the last inequality, that aER (20) bERT
max max (s(a,)) M < as the domain compact.
Similarly, one can show the other way inequality, resulting in the following.
(U(s,, U km)) X(2M)9-1 (10 8m) Ja (lo bm))
(26)
The rate at which IU (sm,tm) tm) - U (so, goes to zero is hence the same that with which cither of the IPM terms goes to zero. For example, if the IPM used for regularization is MMD with normalized kernel, then MMDs (So, 8m)
+ 2log(1/8) with probability at least. 1 (Smola et al. 2007
From the union bound, with probability at least 1 - 5, IU (8ms tm) U (so, to) 2 (2/(2
22(2M)0-1 (V
Thus, O
is the common bound for the rate at which the v
1
LHS well as the MMDs (so. s,) decays to zero.
B.16 Robustness
We show the robustness property of IMMerrguularzeed UOT 13 with the same assumptions on the noise model as used in (Fatras et al. 2021 Lemma 1) for KL-regularized UOT.
Lemma B5. (Robustness) Let so. to E Rt (X). Consider s, pso + (1 p)d, (p E [0, 1)), a distribution perturbed by a Dirac outlier located at some outside of the support of to. Let m(2) c(2, dtdc We have that, 2 (Se, (o) S p (S0, la) + p)m((z)
37

---

Published in Transactions on Machine Learning Research (01/2024)
Proof. Let be the solution of Ug egg,cAi.. (Su, to). Consider ā pi (1 p)ó= @ to. It is easy to see that # pad + p)o= and #2 pT + (1 p)60-
(g 10) - / cz, (I, v) + \ι ( 86) + ( (o) (Using the definition of min)
/ drr v) + ( so) p)7o(62, 8.)) X2 (png (T2, to) + (1 ()o
(:: IPMs jointly convex)
c(,, (I,,))d(T,,) v) + (Arzo(TT, A2To(T2,(o)) /
/ ((x.))dT(E.)) v) / dxz,))dTZ v) + / (1-p)c(2, dtt + + (A su) + (T2, to)) = p Ug 1o) + p)m(z).
@ (o)((. + p ( su) + (22
We note that m(2) finite to (X).
We now present robustness guarantoes with a different noise model.
Corollary B6. We say a measure 4 E (X) is corrupted with E [0, fraction of noise when g (1- p)ge + pqn, where qe the clean measure and On the noisy measure. Let sp,to E RX) be corrupted with p fraction of noise such that - Sn max have that Ugge.,.X to) I Ug.ett( pā(e1 + E2), where fEg(AJNW.
SnlTV TV Sa and ltc tnllr S c2. We
|If
Proof. We use our duality result of Ug from Theorem 4.1 We first upper-bound Cg (8ns (n) which is later used in the proof.
Ug (8n:
max fEG(A)nW.
/ Idan (dsn-- -fathn
saln
/ fd fdn -so)+ + fdse - f fd fd (tn - -to) - fdte
max Seg())nwe
(f fddcc-f fate)
max fEg(A)nW..
fd fsn - -sc) +
max }Eg(X)m....
fd (tn -te)
max
ʃeg())nw.
< 8(1sc 8nlTV + lee (llvv) + (gs.
B(el + =2) + (.G. le).
(27)
We now show the robustness result follows.
f [dso - f fdtt
Ug (g,:2..(50tt00==
max fEG(X)nW:
(1 (1-p)/ [fdsct + P fdsn ((1-p / fate -p ʃ fdtn
max SEC(A)NW
(1- p) (f [ f [Salie + +o Julon |Saki))
max SEG(A)nw.
(1 (1-p) p) ( fdss- 1 ssdto)
[fllsn - fdlt))
< max
max SEg(A)mW.
P
SEG(A)MW.
(1 p) Ug (seyt)) + +p Usg..rd (Snsta)
(1 p) Mo te) Pllg,esd (Sos to) + 3(c + (2)) (Using 27
U....A.A (so) + p3(El + ez).
38

---

Published in Transactions on Machine Learning Research (01/2024)
We note that B
max SEG(A)nW.
max |l/ < 0o. Also, as 3 S min JeW.
max SEGu(A)
Il
max JEW,
115
min
(A, max l1fll(o) rew.
(for a normalized kernel).
B.17
Connections with Spectral Normalized GAN
We comment the applicability of MMD-UOT in generative modelling and draw connections with the Spectral Norm GAN (SN-GAN) (Miyato et al. 2018 formulation.
A popular approach in generative modelling is to define a parametric function ge 2 + X' that takes a noise distribution and generates samples from Pp distribution. We then learn to make Po closer to the real distribution, Py. On formulating this problem with the dual of MMD-UOT derived in Theorem 4.1
we get
/ [aPo - -
min max SEW.
(28)
We note that the above oplimization problem, the critic function or the discriminator f should satisfy Ilf and ||fíl S Å where [lffll denotes the Lipschitz norm under the cost function c. Let the critic func- tion be fw, parametrized using a deep convolution neural network (CNN) with weights W (W,,... Wr),
where L is the depth of the network. Let F be the space of all such CNN models, then Problem (8 can be approximated as follows.
/ fwdde SoALP-ooo /
f fwdde
min www.flllwlllsss.ll.
max
(29)
The constraint liflle I 1 is popularly handled using a penalty on the gradient, f/ (Gulrajani et al. 2017). The raistraint on the RKHS norm, I|]] is more challenging for an arbitrary neural network. Thus, we follow the approximations proposed in (Bietti et al.) 2019). (Bietti et al., (019) use the result derived in (Bietti & Mairal 2017] that constructs a kernel whose RKHS contains CNN, J, with the same architecture and parameters as but with activations that smooth approximations of ReLU, With this approximation, (Bietti et al. 2019) shows tractable bounds on the RKHS norm. We consider their upper
bound based on spectral normalization of the weights in fw. With this, Problem 29 can be approximated with the following.
/ fwdPs - / fwdP, + Pill + P2 fwef
min max
W.l.
(30)
where I|.llsp denotes the spectral norm and P1, P2 O. Formulations like 300 have been successfully applied variants of Spectral Normalized GAN (SN-GAN). This shows the utility of MMMD-regularized UOT in generative modelling.
B.18 Comparison with WAE
The OT problem in WAE (RHS in Theorem 1 in (Tosstikhin et al. 2018 using notation is:
/ (zz G())) d 2), s.t. T. Px. T2 P2,
min (E
(31)
where X' Z are the input and latent spaces, G is the decoder, and Px, Pz are the probability measures corresponding to the underlying distribution generating the given training set and the latent prior (e.g.. Gaussian).
Tolstikhin al. 2018) employs one-sided regularization. More specifically, (Tolstikhin et al. 2018 equ. ()) in our notation is:
/ (Ix dr(x,2) + A2MMD Pz), s.t. T1 Px.
min ER (xx2))
(32)
39

---

Published in Transactions on Machine Learning Research (01/2024)
However, in work, the proposed MMD-UOT formulation corresponding to 31 reads as:
/ c(r, G(z)) dT(x,2) + M. Px) + A2MMD (T2, Pz). min TER (X(
(33)
It. is easy see that the WAE formulation 332 a special case of our MMD-UOT formulation
Indeed,
as A1 co, both formulations are the same.
The theoretical advantages of MMD-UOT over WAE are that MMD-UOT induces new family of metrics and can be efficiently estimated from samples at rate O( (mm whereas WAF is not expected to induce metric as the symmetry is broken. Also, WAE expected to be cursed with dimensions in terms of estimation,
marginal is exactly matched, similar to unregularized OT.
We now present the details of estimating (33 in the context of VAEs. The transport plan и is factorized as <(x =) = कi (=|æ), where ((2/z) is the encoder. For the sake of fair comparison, we choose this encoder and the decoder, G, to be exactly the same as that in (Tolstikhin et al. 2018). Since T (z) is not modelled by WAE, we fall back to the default parametrization our paper of distributions supported over the training points. More specifically, D f.i...., Fmm is the training set (sampled from Px), then our formulation reads as:
min A Σ aj
/ c(2a, c(x), G(2)) dr(2|²;) MMD a, m 111
1
X2MMD (ant())as) Pz
(34)
where is the gram-matrix over the training set D. We solve variables can employ accelerated gradient steps.
using SGD, where the block over the a 34
C Experimental Details and Additional Results
We present more experimental details and additional results in this section. We have followed standard practices to ensure reproducibility. We will open-source the codes to reproduce all our experiments upon acceptance the paper.
C.1 Synthetic Experiments
We present more details for the experiments in Section 5.1 along with additional experimental results.
We use squared-Fuclliden cost the ground metric. We take points [1,2,-.. ,50] and consider Gaussian distribution over them with mean, and standard deviation (15, 5) and (35, 3). respectively. The hyperparameters for MMD-UOT are À as 100 and σ² in the RBF kernel (k(z, y) exp
Transport Plan and Barycenter
v==-"l"") as 1. The hyperparameters for eKL-UOT are À and as 1.
20
For the barycenter experiment, we take points [1, .... 100] and consider Gaussian distribution over them with mean, and standard deviation as (20, 5) and (60, 8), respectively. The hyperparameters for MMD-UOV are À as 100 and 6² in the RBF kernel as 10. The hyperparameters for <KI-UOT are À as 100 and 10-3
Visualizing the Level Sets For all OT variants squared-Fuclidens is used as a ground metric. For the level set with MMD, RBF kernel is used with o² as For MMD-UOT, À is and RBF kernel is used with ²² as For plotting the level set contours, lines are used for all methods.
Computation Time The source and target measures are Uniform distributions from which we sample 5,000 points. The dimensionality of the data is 5. The experiment is donc with hyper-paraneters as squared-Euelidean distance, squared-MMD regularization with RBF kernel, sigma as 1 and lambda as 0.1. eKL-UOT's entropic regularization coefficient is 0.01, and lambda is 1. We choose entropic regularization coefficient from the set {le 3, le 2, le - 1} and lambda from the set {le 2, 2. 1, This hyper- parameter resulted in the fastest convergence. This experiment was done on an AVIDIA-RTX 2080 GPU.
40

---

Published in Transactions on Machine Learning Research (01/2024)
1.0
KL UOT with Sinkharn
Prociem (9) with APCII)
0.8
0.212
Obeecive
0.6
0.210
0.258
0.234
0.232
0.3
Time
Time (5)
Figure Computation time: Convergence plots with m 5000 for the case of the same source and target measures where the optimal objective is expected be 0. Left: MMD-UOT Problem (T)p solved with accelerated projected gradient descent. Right: eKI-UOT's convergence plot is shown separately. We observe that eKL-UOT's objective plateaus in 0.3 seconds. We note that our convergence to the optimal objective is faster than that of eKI-UOT.
MMD-UOT
6x 10
KL-UOT
Objective
4 10
3 10
2 10
103
10
m
Figure 5: Sample efficiency: Log-log plot of optimal objective vs number of samples. The optimal objective values of MMD-UOT and eKL-UOT formulation are shown as the number of samples increases. The data lies in 10 dimensions, and the source and target measures are both Uniform. MMD-UOT can be seen to have a better rate of convergence.
Sample Complexity In Theorem 4.10 in the main paper, proved attractive sample complexity of O mm for our sample-based estimators. In this section, we present a synthetic experiment to show that the convergence of MMD-UOT's metric towards the true value faster than that eKL-UOT. We sample 10-dimensional sources and target samples from Uniform sources and target marginals, respectively. As the marginals are equal, the metrics over measures should converge to as the number of samples increases. We repeat the experiment with increasing number of samples. We use Equared--Enliddan cost. For eKL-UOT, À 1,e = 1e - 2. For MMD-UOT, X = 1 and RBF kernel with we plot MMD-UOT's objective and the square root of the eKL-UOT objective on increasing the number of samples. It can be seen from the plot that the MMD-UOT achieves a better rate of convergence compared to eKL-UOT.
1 is used. In Figure 5
Effect of Regularization In 7 we visualize matching the marginals of MMD-UOT- optimal transport plan. We show the results with both RBF kernel k(r, v) = exp
-1 and the IMQ kernel 2+10
k(æ, y) = (10 6 + ||ا - y 22) 0.5 As we increase X, the matching becomes better for unnormalized measures, and the marginals exactly match the given measures when the mcasures are normalized. We have also shown the unbalanced case results with KL-UOT As the POT library (Flamary et al. 2021 doesn't allow including simplex constraint for KI-UOT, we do not show this.
41

---

Published in Transactions on Machine Learning Research (01/2024)
BF
MMD-UOT with IMQ kermel
MMD-UOT with RBF kernel
64
24
32
KL-UOT
Figure 6: (With unnormalized measures) Visualizing the marginals of transport plans learnt by MMD-UOT and KL-UOT, on increasing X.
640
MMD-UOT with IMQ kemee
4000
MMD-UOT with RBF kemnel
Figure (With normalized measures) Visualizing the marginals of MMD-UOT (solved with simplex con- straints) plan on increasing A. We do not show KL-UOT here the Sinkhorn algorithm for solving KL-UOT in the POT library (Flamary el al. 2021 does not incorporate the Simplex constraints on the transport
plan.
C.2 Two-sample Test
Following Liu et al. 2020), we repeat the experiment times, and in each trial, we randomly sample a validation subset and test subset of size N from the given real and fake MNIST datascts. We run the two-sample experiment for type-Il error on the test set for given trial using the hyperparameters tuned for N 100 for each trial. The hyperparameters
chosen for that trial. The hyperparameters
42

---

Published in Transactions on Machine Learning Research (01/2024)
Table 8: Test power (higher is better) for the task of CTFAR-10.1 vs CIFAR 10. The proposed MMD-UOT
method achieves the best results.
ME
SCF
C2ST-S
C2ST-L
MMD
eKL-UOT
MMD-UOT
0.588
0,171
0.452
0.529
0.316
0.132
0.643
for a given trial were chosen based on the average empirical test power (higher is better) over that trial's validation dataset.
We use squared- distance for MMD-UOT and eKL-UOT formulations. RBF kernel, k(x,y)
exp (- -/==""") is used for MMD and for MMD-UOT formulation. The hyperparameters are chosen from the following set. For the MMD-UOT and MMD, a was chosen from {median, 40, 60, 80, 100} where the median is the median-heuristic |Gretton et al. 2012). For the MMD-UOT an eKL-UOT, A is chosen from {0.1, 1, 10). For eKL-UOT, e was chosen from {1, 10-1, 10-2,100- 10-3, 10-4). Based on validation, a as the median is chosen for MMD at all trials. For eKL-UOT, the best hyperparameters (A, are (10, 0.001) for trial number 3, (0.1, 0.1) for trial number 10 and (1, 0.1) for the remaining the trials. For MMD-UOT, the best hyperparameters (X, are (0.1, 60) for trial number 9 and (1, median") for the remaining 9 trials.
Additional Results Following (Liu et al. 2020 we consider the task of verifying that the datasets CIFAR-10 (Krizhevsky 2009) and CIFAR-10.1 |Recht et al. 2018) are statistically different. We follow the same experimental setup as given in (Liu et al. dataset, and the test is on 1,031 images. The experiment repeated 10 times, and the average test power is compared with the results shown in (Liu et al. 2020) with the popular baselines: ME (Chwialkowski et al., 2015 Jitkrittum et al. 2016), SCF (Chwialkowski ct al. 2015 Jitkrittum et al.] 2016), C2ST-S (Lopez-Paz & Oquab (017)), C2ST-L. (Cheng & Cloninger] (2019). We repeat the experiment following the same setup for the MMD and eKL-UOT baselines. The chosen hyperparameters (A,e) for the 10 different experimen- tal runs UOL-UOT are (0.1,0.1) (1, 0.1), (1, 0.1,, (1,0.
2020). The training is done on 1,000 images from each
(1, (1,0,, (1,0.1), (0.1, 0.1,, (1, 0.1). (1,0.1) The chosen (A,a²) for the 10 different experimental runs of MMD-UOT are
and (1,0.1). (0.1, median), (1, 60), (10, 100), (0),, (0.1,40), (0.1, 40), (0.1, 40), (1, median), (0.1,80) and (1, 40). Table g
shows that the proposed MMD-UÓT obtains the highest test power.
C.3 Single-Cell RNA sequencing
scRNA-seq helps us understand how the expression profile of the cells changes over stages (Schiebinger let al., 2019). A population of cells is represented as a measure of the gene expression space, and as they grow brow/divide, divide/dia and the measure evolves over time. While scRNA-seq records such a measure at a time stamp, does so by destroying the cells (Schiebinger et al., 2019). Thus, is impossible to monitor how the cell population evolves continuously over time. In fact, only a few measurements at discrete timesteps are generally laken due to the cost involved.
We perform experiments on the Embryoid Body (EB) single-cell dataset (Moon et al. 2019 The Embryoid Body datasot comprises data at 5 timesteps with sample sizes as 2381, 4163, 3278, 3665 and 3332, respectively.
The MMD barycenter interpolating between measures 60, has the closed form solution as } (s0 + 1o). For evaluating the performance at timestep t,, we select the hyperparameters based on the task of predicting for
The À hyperparameter for the validation of {ti, f2, ts} 4.. We use IMQ kernel k(z,8) v) = [++ [1+1=-//)²)) -0.5
K²
MMD-UOT is chosen from {0.1, 1, 10} and K²2 is chosen from {le - 4, le - 3, 1e - 2, le - 1, median), where median denotes the median of [0.5]|æ vll2Va, y E D s.l. x A y} over the training dataset (D). The chosen
(A, K2) for timesteps t1, t2, t3 are (1, 0.1), (1, median) and (1, median), respectively. The X hyperparameter for the validation of <KI-UOT chosen from (0.1, 1, 103 and c is chosen from {le 5, le 4, le 3, 3.1e The chosen (X. for timesteps t,, t2, t3 are (10, 0.01), (1, 0.1) and (1, 0.1) respectively. In Table g we compare against additional OT-based methods Wi, W2, cOT.
2. le 1e-1). 1}.
43

---

Published in Transactions on Machine Learning Research (01/2024)
Table 9: Additional OT-based baselines for two-sample test: Average Test Power (between and 1; higher is better) on MNIST. MMD-UOT obtains the highest average test power at all timesteps even with the additional baselines.
N
W1
W2
cOT
MMD-UOT
100
0.111
0.099
0.108
0.154
200
0.232
0.207
0.191
0.333
300
0.339
0.309
0.244
0.588
400
0.482
0.452
0.318
0.762
500
0.596
0.557
0.356
0.873
1000
0.805
0.773
0.508
0.909
KL-UOT
source
target
-300
-90
50
310
MMD UOT
100
source target
-100
Figure 8: (Best viewed in color) The t-SNE plots of the source and target embeddings learnt for the M- MNIST to USPS domain adaptation task. Different cluster colors imply different classes. The quality of the learnt representations can be judged based on the separation between clusters. The clusters obtained by MMD-UOT seem better separated (for example, the red and the cyan-colored clusters).
C.4 Domain Adaptation in JUMBOT framework
The experiments are performed with the same seed as used by JUMBOT. For the experiment on the Digits dataset, the chosen
for MMD-UOT are K2 in the IMQ kernel k(z,y)
[1+
(1+12=-91?
0.5
K
10-2 and À 100. In Figure g UOT and cKL-UOT-based loss. The clusters formed with the proposed MMD-UOT seem better scparated (for example, the red and the cyan-colored clusters). For the experiment on the Office-Home dataset, the
we also compare the t-SNE plot of the embeddings learnt with the MMD-
chosen hyperparameters for MMD-UOT are (X = 100, IMQ kernel with K² = 0.1). For the VisDA-2017
dataset, the chosen hyperparameters for MMD-UOT are (x (>=11 1, IMQ kernel with K²2 as 10).
For the validation phase on the Digits and the Office-Home datasets, we choose À from the set {1, 10, 100) and 2 from the set (0.01, 0.1, 10, 100, median]. For the validation phase on VisDA, we choose À from the set. (1, 10, 100) and K² from the set {0.1, 10, 100).
44

---

Published in Transactions on Machine Learning Research (01/2024)
Original image
PLOT
Proposed MMD-UOT-based Prompt Learning
Figure 9: The attention maps corresponding to each of the four prompts for the baseline (PLOT) and the proposed method. The prompts learnt using the proposed MMD-UOT capture diverse attributes for identifying the cat (Oxford-Pets dataset): lower body, upper body, image background and the area near the mouth.
Original image
PLOT
Proposed MMD-UOT-based Prompt Learning
The attention maps corresponding to each of the 4 prompts for the baseline (PLOT) and the proposed method. The prompts learnt using the proposed MMD-UOT capture diverse attributes for identi- fying the dog (Oxford-Pets dataset): the forehcad and the nose, the right portion the face, the head along with the left portion of the face, and the ear.
Figure 10:
C.5 Prompt Learning
Let F = (fm IM } denote the set of visual features for given image and G, {gnlN | text ual prompt features for class r. Following the setup in the PLOT baseline, an OT distance is computed between empirical measures over 49 image features and textual prompt features, taking cosine similarity cost. Let (x, denote the OT distance between the visual features image x and prompt features of exp(((1-dor((..)//))
} denote the set of
class The prediction probability is given by p(y = r|x) Σ exp (1041110rr) (o0r)oooorrr)))))) where T denotes the
total no. of classes and the temperature softmax. The textual prompt embeddings are then optimizod with the poss-entropyy loss. Additional results on Oxford-Pets (Parkhi et al. 2012) and UCF101 (Soomro
et al. 2012 datasets shown in Table 11
45

---

Published in Transactions on Machine Learning Research (01/2024)
Table 10: Hyperparameters (kernel type, kernel hyperparameter, X) for the prompt. learning experiment.
Dataset
2
4
16
EuroSAT
(imq2, 10 500)
(imql, 10 10
(imql, 10 500)
(imql, 10 500)
(rbf, 1, 500)
DTD
[imql, 10 10-2, 10)
(rbf, 100, 100)
(imq2, 10 10)
(rbf, 10 10)
(rbf, 0.1, 1)
Oxford-Pets
(imq2, 0.01, 500)
(rbf, 10-
(imq. 1, 10)
(imql, 0.1, 10)
(imql, 0.01, 1)
UCF101
(rbf, 1, 100)
(inq2, 10, 100)
(rbf, 0.01, 1000)
(rbf, 10-1, 10)
(rbf, 100, 109)
Table 11: Additional Prompt Learning results. Average and standard deviation (over 3 runs) of accuracy
(higher is better) on the k-shot classification task, shown for different values of shots (k) in the state-of- the-art PLOT framework. The proposed method replaces OT with MMD-UOT in PLOT, keeping all other hyperparameters the same. The results of PLOT are taken from their paper (Chen al. 2023
Dataset
Method
1
2
4
8
16
PLOT
54.05 + 5.95
64.21
1.90
72.36
2.29
78.15 + 2.65
82.23 + 0.91
EuroSAT
Proposed
58.47 = 1.37
66.0 = 0.93
71.97 + 2.21
79.03 = 1.91
83.23 + 0.24
PLOT
46.55 + 2.62
51.24 = 1.95
56.03 + 0.43
61.70 + 0.35
65.60 + 0.82
DTD
Proposed
17.2711.46
51.0+1.71
56.40-
63.17-1.69
65.90 +-t 0.29
PLOT
87.49 + 0.57
86.64
0.63
88.63 + 0.26
87.39
0.74
87.21 = 0.40
Oxford-Pets
Proposed
87.60 + 0.65
87.47
1.04
88.77 + 0.46
87.23
0.34
88.27 + 0.29
PLOT
64.53 +: 0.70
66.83 + 0.43
69.60 + 0.67
74.45 + 0.50
77.26 + 0.64
UCF101
Proposed
64.2 + 0.73
67.47 = 0.82
70.87 0.48
74.87 + 0.33
77.27 0.26
PLOT
63.16
67.23
71.66
75.42
78.08
Aug acc.
Proposed
64.38
67.98
72.00
76.08
78.67
Following the PLOT baseline, we the last-epoch model. The authors empirically found that learning 4 prompts with the PLOT method gave the best results. In our experiments, keep the number of prompts and the other neural network hyperparameters fixed. We only choose À and the kernel hyperparameters for prompt learning using MMD-UOT. For this experiment, we also validate the kernel type. Besides RBF,
consider two kernels belonging to the IMQ family: k(z, ( [++m=-y1)) k(z, (K2 + || -0.5 (referred to as imq2). We choose À from {10, 100, 500, 1000} and kernel hyperparameter (K2 or σ²) from {le 3, le 2, le 1,1, le 2, le 3). The chosen hyperparameters are included in Table 10
(referred to as imql) and
K²
46


Reviews:
[{"review_id": "2UScZ4Fabp", "reviewer": "Reviewer_uqvE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n\n- The authors identify an interesting problem on which little has been done, namely the theoretical foundations of UOT with MMD-based regularization.\n- Questions such as the structure of the underlying metric and statistical consistency of practical estimators are significant, since methods like UOT continue to be employed in real-world generative modeling tasks, and the authors do a nice job at explaining some of this story and providing rigorous proofs.\n- I find their duality and interpolation results very interesting from the perspective of traditional optimal transport. It really explains what the metric induced by MMD-regularized UOT is capturing.\n- The paper is generally well-written, making it easy to understand their main results.\n\nWeaknesses:\n\n- The main weakness is that the assumption of Theorem 4.10 seems quite unrealistic and thus precludes the application of their work to many interesting settings. In particular, the authors assume that the function $\\eta(x,y) := \\frac{\\pi^{\\ast}(x,y)}{s_0(x) t_0(y)}$ belongs to $\\mathcal{H}_k \\otimes \\mathcal{H}_k$, where $s_0$ and $t_0$ are the two measures we wish to compare, $\\pi^{\\ast}(x,y)$ is the optimal transport plan between $s_0$ and $t_0$, and $\\mathcal{H}_k$ is the RKHS. They also assume that the cost function $c(x,y)$ used to define their optimal transport plan belongs to $\\mathcal{H}_k \\otimes \\mathcal{H}_k$. In addition, the error of their estimator has a constant of $\\|\\eta\\|_k \\cdot \\|c\\|_k$, the RKHS norms of $\\eta$ and $c$.\n\n- My concern is that this assumption assumes far too much regularity on the target measures and cost function, since the RKHS of a characteristic kernel can be quite small (evidenced, for example, by the $O(n^{-1/2})$ rate for the uniform law of large numbers that the paper cites). It does not allow for non-absolute continuity of the distributions, which is a common assumption in machine learning. The authors state that the assumption is mild because the RKHS is dense in the space of continuous functions. However, they do not provide any quantitative approximation results, i.e., they do not say, for given $f \\in C(\\mathcal{X}), $\\epsilon > 0$, what the RKHS norm of a function $f_{\\epsilon} \\in \\mathcal{H}^k$ that satisfies $\\|f-f_{\\epsilon}\\|_{\\infty} < \\epsilon$ will be. Also, it is not clear why the cost would belong to the RKHS.\n\nI think the paper would be much improved if the authors:\n\n- Were able to relax the assumptions on $s_0$, $t_0$, and $c$ or weaken the dependence of the sample complexity bound on the RKHS norms.\n- Put more classical/realistic assumptions on the target measures and consider approximation as an additional source of error.\n\nIf the authors insist that it is a ‘mild’ assumption, then they should provide at least one test example where the assumptions are satisfied. Otherwise, it should be mentioned that the assumption is restrictive, and future improvements should be discussed in the conclusion. \n\nI also have some minor issues with the presentation. First, the notions of characteristic and universal kernels are used throughout the paper but never defined. It would improve readability to have a brief discussion of what kernels are included/excluded from the characteristic assumption, even if it is delayed to the appendix. Also, I think you should mention what kernels you use for your experiments in the main text.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors study unbalanced optimal transport as a tool to compare probability measures, which is a central task in machine learning. Unbalanced optimal transport (UOT) is a variant of the standard Kantorovich problem, wherein the hard constraints on the transport plans are replaced with soft regularization terms that impose the marginals of the transport plan to be close to the source and target measures. The authors study the analytical and statistical properties of UOT when the regularization terms are defined using the maximum-mean discrepancy (MMD) metric on probability measures corresponding to a reproducing kernel Hilbert space (RKHS). On the analysis side, the authors prove a duality formula for MMD-regularized UOT that is analogous to the duality for the traditional Kantorovich problem in optimal transport. Under certain assumptions on the kernel, they also prove that MMD-regularized UOT is a genuine metric on probability measures, and that it interpolates between the Kantorovich metric and a rescaled version of the MMD metric, depending on the regularization parameter. On the statistics side, the authors propose an estimator for MMD-regularized UOT when the measures are only accessible through data, which can be computed efficiently using convex optimization methods. They prove that their estimator converges to the true metric in the infinite-data limit with a dimension-independent convergence rate. The paper concludes with the discussion of several synthetic experiments, including the computation of barycenters.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I suggest the authors make changes by addressing the following questions. \n\n-\tWhat are the benefits of employing MMD-UOT in practice as opposed to just using MMD? It seems like they have similar statistical properties, and computing MMD does not require one to solve an optimization problem.\n-\tI wonder if it is worth considering MMD-regularized UOT with kernels that are not characteristic (e.g., polynomial kernels), for the reasons described earlier. Perhaps you would get a worse statistical convergence rate, but the result might apply to a wider variety of measures. In addition, you would probably lose the metric properties, although it seems like for many applications it doesn’t really matter whether you are using a genuine metric. Have you considered this?\n-\tIs the IMQ kernel characteristic?\n-\tIn the introduction, you write “Since MMD-UOT can approximate Wasserstein arbitrarily closely (as the regularization hyperparameter goes to $\\infty$), our result can also be understood as a way of alleviating the curse of dimensionality problem in Wasserstein.” Do you have any quantitative estimates for the rate of convergence of MMD-UOT to the Wasserstein metric as the regularization parameter increases? If this rate suffers from a curse of dimensionality, then perhaps it is not quite true that using MMD-UOT breaks the curse of dimensionality for the Wasserstein metric.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "JxTP6fok5d", "reviewer": "Reviewer_LYuh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- This paper makes a strong contribution of considering the popular MMD for constraint regularization which is a very practical choice.\n- The dual form derived is very interesting as it translates to constraining the LP dual (Kantorovich-Rubenstein duality) in Optimal Transport to the function classes being in the RKHS and in the case of IPMs the weighted generator class.\n- The algorithm presented is very practical and benefits from MMD.\n\nWeaknesses:\n- The theoretical result appears to be minimizing over unsigned Borel measures as opposed to probability measures and therefore only a generalization of this variant of OT.\n- When using Sion's minimax theorems, compactness of the set is required in one of the sets. This makes the IPM result slightly limited.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this work, a study of relaxed optimal transport is conducted where the marginal constraints are distributed into the objective in the form of penalties. While the standard choice to penalize these constraints is the phi-divergence (or the f-divergence as commonly known in the literature), this work studies the case of selecting the kernel-Maximum Mean Discrepancy (MMD). It is shown that the optimization objective admits a dual formulation which has conceptual appeal. Additionally, it is shown how this formulation can be computed with competitive estimation rates of O(m^{-1/2}). Experiments on various applications of machine learning are then demonstrated.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think it would be better if the authors mention that the optimization relaxed considers not only probability measures but also unsigned Borel measures which is a much larger set. While this is implicit in the study of \"unbalanced\" optimal transport, it would be better if this more explicitly mentioned. I still think the paper should be accepted in its current state without this. Additionally, I think it would be better if the authors motivate unbalanced OT or how signed (unnormalized) Borel measures appear naturally in machine learning applications. \n\nPage 2: \"Such efficient estimators are particularly useful in machine learning applications, where typically only samples from the underlying measures are available\"\nIt would be better if you can specify which machine learning applications are useful here.\n\nI think you should include the objective of the Wasserstein Autoencoder [1] in the discussion of related work since in their work the \"WAE-MMD\" closely resembles the study where the marginal constraint is penalized with MMD. In particular, marginal constraint being penalized using MMD (in exactly WAE-MMD) was studied in [2] where MMD regularization leads to a variational objective involving RKHS functions  (see Equation (25) ) which is highly relevant to the dual form derived here.\n\nIn the proof of Theorem 4.1, I would request the authors to explain the details of using Sion's minimax Theorem. In particular, it is required that at least one of the spaces (either $\\mathcal{R}^{+}(\\mathcal{X} \\times \\mathcal{X})$ or $\\mathcal{G}$) are compact. I would recommend the authors to clearly state this fact so the compactness assumption of $\\mathcal{G}$ can be better understood.\n\nIf $\\mathcal{X} \\times \\mathcal{X}$ is assumed to be compact then one can actually rely on the compactness of $\\mathcal{R}^{+}(\\mathcal{X} \\times \\mathcal{X})$ to make Sion's theorem work and a slightly different result would be present. I would encourage the authors to add such a discussion.\n\n\n[1] Tolstikhin, I., Bousquet, O., Gelly, S., & Schoelkopf, B. (2017). Wasserstein auto-encoders. arXiv preprint arXiv:1711.01558.\n\n[2] Husain, H., & Knoblauch, J. (2022, March). Adversarial interpretation of Bayesian inference. In International Conference on Algorithmic Learning Theory (pp. 553-572). PMLR.", "broader_impact_concerns": "As this work studies the theory and abstraction of machine learning application, there are no immediate ethical concerns in the work itself.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Mm6HyKtlG7", "reviewer": "Reviewer_rbdr", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1) The proposed method is reasonable and solid, with theoretical supports, and the theory in the submission is proven in details.\n\n2) In the aspect of sample complexity, the proposed MMD-UOT suppress the curve-of-dimensionality issue, which provides a good surrogate of Wasserstein distance and KL-divergence-based UOT in high-dimensional applications.\n\n3) The paper is well-written, with a good logic flow, which is friendly to the readers without much background.\n\nWeaknesses:\n\n1) My main concern is about the experimental part of this work. In particular, the experimental part of this submission is weak to some degree. Firstly, the scale of the datasets is not large enough, which cannot demonstrate the usefulness of the proposed method in large-scale applications. Secondly, in each experiment, the baselines are insufficient --- only the MMD and the KL-UOT are considered. Especially in the generative modeling task, only toy datasets like MNIST are considered, and many OT-based generative models (e.g., WGAN [a] and its IPM extensions) are ignored. \n\n2) Additionally, for the MMD-UOT based barycenter, besides the toy experiment shown in Figure 2, more practical applications and datasets should be considered.\n\n3) In the section of preliminaries, to make the paper more friendly to readers, the authors should show the dual form of Wasserstein distance explicitly.\n\n4) The Wasserstein autoencoder (WAE) in [b] should be added in the literature review and experimental part, which also considers an MMD regularizer when solving an OT problem approximately. \n\n[a] Arjovsky, Martin, Soumith Chintala, and Léon Bottou. \"Wasserstein generative adversarial networks.\" International conference on machine learning. PMLR, 2017.\n\n[b] Tolstikhin, I., et al. \"Wasserstein Auto-Encoders.\" 6th International Conference on Learning Representations (ICLR 2018). OpenReview. net, 2018.\n\nMinors:\n\n1) Figure 1 is meaningless in my opinion, which can be deleted in the revised paper.\n\n2) In the implementation, does the proposed method require the source and target measures to have the same number of samples (i.e., the \"m\" in the paper)? If the answer is No, a footnote should be added to explain the case with different numbers of samples.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this submission, the authors propose a new variant of unbalanced optimal transport (UOT) called MMD-regularized UOT. \nCompared with classic KL-divergence-regularized UOT, the proposed method suppresses the curse-of-dimensionality issue and thus may be more suitable for high-dimensional applications. \nAdditionally, the authors investigated the theoretical property of the MMD-UOT, demonstrating that 1) it is a valid metric, working as a surrogate for both MMD and Wasserstein distance; 2) it achieves an \"interpolation\" of MMD and Wasserstein distance --- changing the weight of the regularizer leads to different metrics.\nIn the aspect of the computation, an accelerated projected gradient descent algorithm is applied to compute the MMD-UOT with good empirical convergence, and the barycenter problem based on the MMD-UOT is considered.\nExperiments on synthetic and real-world data show the rationality of the proposed method.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See the weaknesses mentioned above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Initial reviews praised several aspects of the submissions (theoretical contributions, relevance, scope) and suggested improvements on others (relaxing the assumptions, expanding the numerical experiments). The authors very convincingly addressed all these points in their rebuttals and all three reviewers unanimously recommend acceptance, which I am happy to follow. The paper contributes an interesting strategy to study unbalanced optimal transport problems through regularisation with the MMD. Congratulations on a fine piece of work!
Decision: Accept (poster)