Statistical and Computational Guarantees of Kernel Max-Sliced Wasserstein Distances
Authors: Jie Wang
March Boedihardjo
Yao Xie
Venue: ICML 2025 poster (2025)
Keywords: Optimal transport
finite-sample guarantees
rank-constrained optimization
Primary area: general_machine_learning->kernel_methods
TL;DR: Statistical and computational guarantees for the kernerlized version of max-sliced Wasserstein distances are provided

Abstract: Optimal transport has been very successful for various machine learning tasks; however, it is known to suffer from the curse of dimensionality. Hence, dimensionality reduction is desirable when applied to high-dimensional data with low-dimensional structures. The kernel max-sliced (KMS) Wasserstein distance is developed for this purpose by finding an optimal nonlinear mapping that reduces data into $1$ dimension before computing the Wasserstein distance. However, its theoretical properties have not yet been fully developed. In this paper, we provide sharp finite-sample guarantees under milder technical assumptions compared with state-of-the-art for the KMS $p$-Wasserstein distance between two empirical distributions with $n$ samples for general $p\in[1,\infty)$. Algorithm-wise, we show that computing the KMS $2$-Wasserstein distance is NP-hard, and then we further propose a semidefinite relaxation (SDR) formulation (which can be solved efficiently in polynomial time) and provide a relaxation gap for the obtained solution. We provide numerical examples to demonstrate the good performance of our scheme for high-dimensional two-sample testing.

Statistical and Computational Guarantees of Kernel Max-Sliced Wasserstein Distances
Jie Wang
March Boedihardjo Yao Xie
Abstract
tion [2, 12-14, 77], etc. See [62] for comprehensive reviews on these topics. The sample complexity of the Wasser- stein distance has been an essential building block for OT in statistical inference. It studies the relationship between a population distribution μ and its empirical distribution a Ei=1 6z, with x; ~ fs in terms of the "Wasserstein dis- tance". Unfortunately, the sample size needs be ex- ponentially large in data dimension to achieve accurate enough estimation [22], called the curse of dimensionality issue.
Optimal transport has been very successful for various machine learning tasks; however, it is known to suffer from the curse of dimension- ality. Hence, dimensionality reduction is de- sirable when applied high-dimensional data with low-dimensional structures. The kernel max- sliced (KMS) Wasserstein distance is developed for this purpose by finding an optimal nonlinear mapping that reduces data into dimension be- fore computing the Wasserstein distance. How- ever, its theoretical properties have not yet been fully developed. In this paper, we provide sharp finite- finite-sample sample guarantees under milder technical assumptions compared with state-of-the-an for the KMS p- p-Wasserstein distance between two empirical distributions with n samples for gen- eral p E [1, oo). Algorithm-wise, we show that computing the KMS 2- W-ssssesteein distance NP-hard, and then we further propose a semidefi- nite relaxation (SDR) formulation (which can be solved efficiently in polynomial time) and provide a relaxation gap for the obtained solution. We pro- vide numerical examples to demonstrate the good performance of our scheme for high- lige-diiensional swo-sample testing.
To tackle the challenge of high dimensionality, it is mean- ingful to combine OT with projection operators in low- dimensional spaces. Researchers first attempted to study the sliced Wasserstein distance [10, 11, 17, 38, 39, 54, 56), which computes the average of the Wasserstein distance between two projected distributions using random one- dimensional projections. Since single random projec- tion contains little information to distinguish two high- dimensional distributions, computing the sliced Wasserstein distance requires large number of linear projections. To ad- dress this issue, more recent literature considered the Max- Sliced (MS) Wasserstein distance that seeks the optimal pro- jection direction such that the Wasserstein distance between projected distributions is maximized [16. 45, 47, 60, 74]. Later, Wang et al. [75] modified the MS Wasserstein dis- tance by seeking an optimal nonlinear projection belong- ing to a ball of reproducing kernel Hilbert space (RKHS), which call the Kernel Max-Sliced (KMS) Wasserstein distance. The motivation is that nonlinear projector can be more flexible in capturing the differences between two high-dimensional distributions; it is worth noting that KMS Wasserstein reduces to MS Wasserstein when specifying a dot product kernel.
1. Introduction
Optimal transport (OT) has achieved much success in var- ious areas, such as generative modeling [27, 49, 59, 61],
distributionally robust optimization [24, 25, 73], non- parametric testing [63, 71, 76, 78, 80], domain adapta-
Despite promising applications of the KMS Wasserstein distance, its statistical and computational results have not yet been fully developed. From statistical perspective, Wang et al. [75] built concentration properties of the empiri- cal KMS Wasserstein distance for distributions that satisfy the projection Poincaré inequality and the Poincaré inequal- ity [42], which could be difficult to verify in practice. From computational perspective, the authors therein developed mradient-based algorithm to approximately compute the empirical KMS Wasserstein distance. However, there is no
"School of Artificial Intelligence, The Chinese University of Hong Kong, Shenzhen, Shenzhen, China "School of Data Sci- ence, The Chinese University of Hong Kong, Shenzhen, Shenzhen, China DDepartment of Mathematics, Michigan State University, East Lansing, USA* "School of Industrial and Systems Engineering, Georgia Institute Technology, Atlanta, USA. Correspondence to: Yao Xie <yao.xie@isye cew.oiieeiiygggtcc........ gatech.edu>.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Guarantees of KMS Wasserstein Distances
theoretical guarantee on the quality of the local optimum solution obtained. In numerical experiments, the quality of the solution obtained is highly sensitive to the initialization.
isting literature. From statistical perspective, existing results on the rate of empirical MS/KMS Wasserstein either dimension-depenndntt suboptimal or require regular- ity assumptions (e.g., log-concavity, Poincaré inequality, projection Bernstein tail condition) on the population dis- tributions [3, 47, 57, 74], except for the very recent litera- ture [9] that provides a sharp, dimension-free rate for MS Wasserstein with data distributions supported on a compact subspace but without regularity assumptions. From the com- putational perspective, there are two main approaches to compute such distances. One is to apply pradient-based al- gorithms to find local optimal solutions or stationary points, see, e.g., [33-35, 46, 75]. Unfortunately. due to the highly non-convex nature of the optimization problem, the quality of the estimated solution is unstable and highly depends on the choice of initial guess. The other is to consider solv- ing its SDR instead [60]. yet theoretical guaranteess on the solution from convex relaxation are missing. Inspired by existing reference [4, 18, 44, 58] that studied the rank bound of SDR for various applications, we adopt their proof tech- niques to provide similar guarantees for computing KMS in Theorem 4.5. Besides, all listed references add entropic regularization to the inner OT problem and solve the regu- larized version of MS/KMS Wasserstein distances instead, while the gap between the solutions from regularized and original problems could be non-negligible
To address the aforementioned limitations, this paper pro- vides new statistical and computational guarantees for the KMS Wasserstein distance. Our key contributions are sum- marized as follows.
We provide a non-asymptotic estimate on the KMS p- Wasserstein distance between two empirical distri- butions based on n samples, referred to as the finite- sample guarantees. Our result shows that when the samples are drawn from identical populations, the rate of convergence is n -1/(2p), which is dimension free and optimal in the worst case scenario.
We analyze the computation of KMS 2- W-sssserstein distance between two empirical distributions based on samples. First, we show that computing this distance exactly is NP-hard. Consequently, we are prompted to propose semidefinite relaxation (SDR) as an approximate heuristic with various guarantees.
We develop an efficient first-order method with biased gradient oracles to solve the SDR, the com- plexity of which for finding 6-optimal solution is Ô (n²8-2 max(n, d-))). In comparison, the complexity of the interior point method for solv- ing SDR is (m55 [5].
Our proposed KMS Wasserstein distance is also closely related to [83], which considers the Wasserstein distance be- tween pushforward measures o and 2g V via an implic- itly defined kernel map .(-). We will show that our metric can be viewed as the max-sliced version of this setting (Re- mark 2.6). A key distinction is that our method enjoys sharp statistical convergence rates, which are difficult to obtain in Zhang et al. [83] due to the curse of dimensionality of Wasserstein.
We derive theoretical guarantees for the optimal solutions from the SDR. We show that there exists an optimal solution from SDR that is at most rank- k, where k - 1 + 22 + 9/4 - 3/2], whereas computing the KMS distance exactly requires a rank-1 solution. We also provide a correspond- ing rank reduction algorithm designed to identify such low-rank solutions from the pool of optimal solutions of SDR.
2. Backkround
We exemplify our theoretical results in non- carammamettic two-sample testing, human activity detection, and gen- erative modeling. Our numerical results showcase the stable performance and quick computational time of our SDR formulation, as well as the desired sample complexity rate of the empirical KMS Wasserstein dis- tance.
We first introduce the definition of Wasserstein and KMS Wasserstein distances below.
Definition 2.1 2. (Wasserstein Distance). Let p E [1, oo). Given a normed space (V. Il II), the p-Wasserstein distance between two probability measures μr, v on V is defined as
/ ||:0 vlll du((,,))
1/p
wp(n,
min (ET((.))
Literature Review
where I(u, v) denotes the set of all probability measures on V x V with marginal distributions μ and v.
In the following, we compare our work with the most closely related literature, and defer the detailed comparision of KMS Wasserstein distance with other variants of OT di- vergences in Appendix A.
Definition 2.2 (RKHS). Consider symmetric and positive definite kernel K : B x B + IR, where B C RRd, Given such a kernel, there exists a unique Hilbert space H. called RKHS, associated with the reproducing kernel K. Denote by K2 the kernel section x E B defined by K₂(y) = K K(x,, v), Vy e
The study on the statistical and computational results of MS and KMS Wasserstein distances is popular in the ex-
2

---

Guarantees of KMS Wasserstein Distances
B. Any function f H satisfies the reproducing property
f(x) (f,Ka Ka)H, Vx B. For y E B, it holds that kx,,y)- = (Ka, Ky)H.
Damnet Y
Delame?
Definition 2.3 (KMS Wasserstein Distance). Let p E [1, o.). Given two distributions and v, define the y-KMS Wasserstein distance as
Dones
Dumon
Scatter plot of cir- (b) Density plot (c) Density plot using
max (.M. SEH, I1
ele dataset
ing MS Wasserstein KMS Wasserstein
where [y and f# are the pushforward measures of j and y by f: + R. respectively.
075
6.75
In particular, for dot product kernel K(x,3) 3) = x zTy. the RKHS H = =f f(x) = x²x 3, 38 E ]. In this case, the KMS Wasserstein distance reduces to the max-sliced Wasserstein distance [16]. A more flexible choice is the Gaussian kernel K (z, y) = - vll2), where
120 0U E20 COO
-18
9.8 3J5 C30 0U5 150
(d) Plot of f(x) = ||||2, E [o. 1
(e) Plot of estimated projector using KMS Wasserstein
> 0 denotes the temperature hyper-parameet In this case, the function class H. satisfies the universal property as is dense in the continuous function class with respect to the oo-type functional norm. is easy to see the following theorem holds.
Figure I. Results on 2-dimensional toy example.
light of this, for two probability measures je and v on H. we define the MS p- D-sssserstein distance
Theorem 2.4 (Metric Property of KMS). K.MS sat-
isfies the triangle inequality. If additionally assuming H is a universal RKHS, K.MS,(U, v) 0 if and only if v.
MS =
We Wo (()))HH. (uy) (2) JEH: I1
sup
We present toy example that highlights the flexibility of the KMS Wasserstein distance. Fa) displays scatter plot of the circle dataset, which consists of two groups of samples distributed along inner and outer circles, perturbed by Gaussian noise. Since the data exhibit a nonlinear structure, distinguishing these groups using a lin- ear projection is challenging. As shown Fig. 1(b), the den- sity plot of the projected samples using the MS Wasserstein distance with linear projector is not sufficiently discrim- inative. In contrast, Fig. I(c) demonstrates that the KMS Wasserstein distance is better suited for distinguishing these two groups Intuitively, the optimal nonlinear projector should take the form (x) x ||=|1 for some scalar c > O, as illustrated in Fig. (d). We plot the nonlinear projector by computing the empirical KMS Wasserstein distance, shown in Fig. (e),, which closely resembles the circular landscape depicted in Fig. 1(d). This result demonstrates that the KMS Wasserstein distance provides data-driven, non-parametric nonlinear projector capable of effectively distinguishing distinct data groups.
Example 2.5.
where (uy) f denotes the pushforward measure of p by the map u i.e., if μ is the distribution of random element X of H. then (uy) (1))2/1 2/11 is the distribution of the random variable (X)= = (J, x) and (uz) ) V is defined likewise. In the following, we show that the KMS Wasserstein distance in Definition 2.3 can be reformulated as the MS Wasserstein distance between two distributions on (infiniteddimensionaa) Hilbert space.
Remark 2.6 (Reformulation of KMS Wasserstein). By the reproducing property, we can see that fxx) (S, K,) u (f(z)) which implies f uy t. As a consequence,
K.M,((,,))
Wp (U)) (D (uF) (o+v) (3)
sup SEH IIfIISSI
=MS, (THHi, $4V).
In other words, the KMS Wasserstein distance first maps data points into the infiite-dimensional Hilbert space H. through the canonical feature map T and next finds the linear projector to maximally distinguish data from two populations. Compared with the traditional MS Wasserstein distance [16] that performs linear projection in RRd KMS Wasserstein distance is a more flexible notion.
Given the RKHS H. let the canonical feature map that embeds data to H as
Φ B H.
x + ((x) = Kg.
(1)
Define the functional uy H + R by uy(g)= (f g for any E H. which can be viewed as a linear projector that maps data from the Hilbert space H to the real line. In
Remark 2.7 (Connections with Kernel PCA). Given data , In on B, denote by f.r the corresponding empirical distribution. Assume 1 Eiejm +(xx) O. since
points Zi....,
3

---

Guarantees of KMS Wasserstein Distances
otherwise one can center those data points as a preprocess- ing step. Kernel PCA [52] is popular tool for nonlinear dimensionality reduction. When seeking the first princi- pal nonlinear projection function f. Mairal and Vert [51]
where C 2 1 is a universal constant. We show the fol-
lowing finite-sample guarantees on the KMS p- p-Wasserstein distance.
Theorem 3.2 Finite-Sample Guarantees). Fix p E (0),
presents the following reformulation of kernel PCA:
level a E (0, 1), and suppose Assumption 3.1 holds.
arg max Var (mz)
(4)
(I) One-Saaple Guarantee) With probability at least 1 a, it holds that K.MS P)) S (n,
SEN: I1
where Var(-) denotes the variance of given probability measure. In comparison, the KMS Wasserstein distance aims to find the optimal nonlinear projection that distin- guishes two populations and replaces the variance objective in (4) with the Wasserstein distance between two projected distributions in (3). Also, kernel PCA is special case of KMS Wasserstein by taking p 2, μ = Mr, 6, in (3).
(II) (Two-Sample Guarantee) With probability at least
1 1- o. it holds that
K.MS Bn) K.M, + A a).
The proof of Theorem 3.2 is provided in Appendix B. The dimension free upper bound A(n, a) = O(n is optimal in the worst case, Indeed, in the one-dimension case B = [0, 1] and K(x,y) = ry. the kernel max-sliced Wasserstein distance K.MS,, coincides with the classical Wasserstein distance Wp. In this case, it is easy to see that if μ - (60 + 61)/2 22 is supported on the two points 0 and 1, the expectation of K.MS( pn,) is of order n 1/2p) [22]. We also compare this bound with other OT divergences in Appendix A.
Notations. Let ( denote the inner product operator. For any positive integer n, denote [n] = (,, n.). Define In as the set
R RAXX
Σ T
1 n
Σ
Fij n
Vi,j e m
(5) Let Conv(P) denote convex hull of the set P. and St denote the set of positive semidefinite matrices of size n x n. We use (.) as a variant of (-) to hide logarithmic factors.
We design a two two-sample test TKMS such that Ho is rejected if K.MS,, (7,, 2n,) > A (nn, a)). By Theorem 3.2, we have the following performance guarantees of TKMS-
3. Statistical Guarantees
Corollary 3.3 (Testing Power of TKMS). Fix a level a e (0.1/2), p E [L, oo), and suppose Assumption 3.1 holds. Then the following result holds:
Suppose samples x" := (arhe(t) and y" := f& he(n are given and follow distributions Jt, V, respectively. Denote by fin and in the corresponding empirical distributions from samples xn and y". In this section, we provide finite-sample guarantee on the p-KMS Wasserstein distance between oh and 2n with E [1, oo). This guarantee can be helpful for KMS Wasserstein distance-based hypothesis testing that has been studied in [75]: Suppose one aims to build a non- con-maraettric test to distinguish two hypotheses Hu μ === v and H1 μ # Thus, it is crucial to control the high-probability upper bound of K.MS, (P,, Dy)) under Ho as it serves as the critical value to determine whether Ho is rejected or not. We first make the following assumption on the kernel.
(I) (Risk): The type-I risk of TRMM is at most a;
(II) (Power): Under H1 : a + v, suppose the sam- ple size n. is sufficiently large such that On == K.M.p(u,)) ( A(n, a) > O, the power of TKMs -.-..nnn//0 -11((p)), where e is a constant is at least 1 - c n depending on A, C, On-
The proof of Corollary 3.3 is provided in Appendix B. It is also noteworthy that the performance of TRMS still has a non-nee neglinglie dependence on the data dimension. Let us follow Ramdas et al. [64] to consider the fair alternative that the discrepancy between μ and v remains constant as the dimension increases. Then, the term @. decreases as the dimension increases. To ensure that the assumption en. >> 0 holds in Corollary 3.3, the sample size should increase to maintain satisfactory testing power when the data dimension increases.
Assumption 3.1. There exists some constant A such that the kemne ((s))) satisfies /K(x, x) S A, Vx E B.
Assumption 3.1 is standard in the literature (see, e.g., [28]), and is quite mild: Gaussian kernel K(I,z) = exp(- || vll3/o²) naturally fits into this assumption. For dot product kernel K K(x, y) = x² xTyy y. if we assume the support B has a finite diameter, this assumption can also be satisfied. Define the critical value
Remark 3.4 (Comparison with Maximum Mean Discrep- ancy (MMD). MMD has been popular kernel-based tool to quantify the discrepancy between two probability mea- sures (see, e.g., [6, 23, 28, 37, 48, 53, 66, 67, 69, 72]), which,
A(n, = 4A(c- C+ 41 log 2) a
2/1/2 1/p
nn 1/(22))
4

---

Guarantees of KMS Wasserstein Distances
for any two probability distributions yr and v, defined as
Cholesky decomposition G-1 1 UUT. Define
MMD(U, =
max JEH. 1f
By Ey[f]
Mi = ((K(0))
(K(1z, 01) - K (yj, 3.1 )E(M) (K (15, 31) - K (zi, 31 )eem
E R
(6)
and the vector M.,- = U M.J By substituting the expres- sion (7) into (KMS) and calculation (see Appendix D), we obtain the exact reformulation of (KMS):
= max SER. 11
(uy) (p))) - (uy) (W
where & denotes the mean of given probability measure 5. The empirical (biased) MMD estimator also exhibits dimension-free finite-sample guarantee as in Theorem 3.2: it decays in the order of (n 1/2), where n: is the number of samples. However, the KMS Wasserstein distance more flexible as it replaces the mean difference objective (6) by the Wasserstein distance, which naturally incorporates the geometry of the sample space and is suitable for hedging against adversarial data perturbations [24].
max WER 1
(W)) min TEE
(9)
Here, we omit taking the square root of the optimal value of the max-min optimization problem for simplicity of pre- sentation. Since Problem (9) a non-convex program, it is natural to question its computational hardness. The follow- ing gives affirmative answer, whose proof is provided in Appendix E.
Theorem 4.2 (NP-hardness). Problem (9) is NP-hard for the worst-case instances of (M.
4. Computing 2-KMS Wasserstein Distance
Let Pi and 7n be two empirical distributions supported on n points, i.e., Th, = Σ. 8g. Fy = Σ, 6g.. where (ix)i, ) are data points in RRd This section focuses on the computation of 2- 2-KMS Wasserstein distance between these two distributions. By Definition 2.3 and monotonicity
The proof idea of Theorem 4.2 is to find an instance of { I that depends on a generic collection of n vec- tors {A,); 1s such that solving (9) is at least as difficult as solving the fair-PCA problem [65] with rank-1 matrices (or fair beamforming problem [68]) and has been proved to be NP-hard [68]. Interestingly, the computational hardness of the MS Wasserstein distance arises from both the data di- mension d and the sample size n, whereas that of the KMS Wasserstein distance arises from the sample size n only.
of square root function, it holds that
XMSS2((n, Dr)
1/2
max SEN. 11
min WEF
Σ (.E[M)]
m=)l((x)) - -((s))
(KMS)
To tackle the computational challenge of solving (9), in the subsequent subsections, we present an SDR formula and propose an efficient first-order algorithm solve it. Next, analyze the computational complexity of our proposed al algorithm and the theoretical guarantees on SDR.
where In, is defined in (5).
Although the outer maximization problem is a ffunctionall op- timization that contains uncountably many parameters, one can apply the representer theorem (see below) to reformulate Problem (KMS) as a fiite-dimennional optimization.
4.1. Semidefinite Relaxation with Efficient Algorithms
Theorem 4.1 (Theorem in [75]). There exists an optimal solution to (KMS), denoted as f, such that for any z,
We observe the simple reformulation of the objective in (9):
Σ Ema((M = Σ Tiis MM >W
n
f(z)
Σ an, K(2
Σ ((,
Inspired by this relation, we use the change of variable approach to optimize the rank-1 matrix S = wwT, i.e., it suffices to consider the equivalent reformulation of (9):
where a = (as,))eln) ay (ays))E(m) are coefficients to be determined.
max {
(F(S): S E 32n Trace(S) rank(S) = 1
Define gram matrix K [(x",>")) (K(a,, (1,),, (r),r,))))ee Je and other gram matrices K(x", """,, K(v", ,,",, K K(y", likewise, then define the concatenation of gram matrics
where F(S) min xef.
Σ n.. (Ma
G K(x",x") x²) (K((""
-K(", K("","")
(10)
2022n
An efficient SDR is to drop the rank-1 constraint to consider the semidefinite program (SDP):
Assume G is positive definite such that it admits the
max F(S), SESm.
In Appendix C. we provide sufficient condition that ensures G positive definite.
where Szn = {SE Trace(S) = =13.
(SDR)
5

---

Guarantees of KMS Wasserstein Distances
Remark 4.3 (Connection with [80D. We highlight that
Algorithm Inexact Mirror Ascent for solving (SDR)
Xie and Xie [80] considered the same SDR heuristic to compute the MS 1-Wasserstein distance. However,
1: Input: Max iterations T, initial guess S1, tolerance c,
constant stepsize 7.
the authors therein apply the interior point method to solve a large-scale SDP, which has expansive complexity
2: for k = 1,..., ,
do
3: Obtain a e-optimal solution (denoted as 7) to (11)
( (oooyoooo) (up to d-aceuracy) [5]. In the following, we present a first-order method that exhibits much smaller complexity ① (n28 3)) in terms of the problem size n (see Theorem 4.4). Besides, theoretical guarantees the solu- tion from SDR have not been explored in [80], and we are the first literature to provide these results.
4: Construct incxact supgradient v(Ss) by (12)
Perform mirror ascent by (13)
end for
Return S1.T + +k= Sk.
which admits the following closed-form update:
The constraint set S2n called the Spectrahedron and ad- mits closed-form Bregman projection. Inspired by this, we propose an inexact mirror ascent algorithm to solve (SDR). Its high-level idea is to iteratively construct an inexact gradi- ent estimator and next perform the mirror ascent on iteration points. By properly balancing the trade-off' between the bias and cost of querying gradient oracles, this type of algorithm guaranteed to find a near-optimal solution [29-32].
5k+1 exp (log Sk + +o((S)),,
3x+1
(13)
Sk+1 (Tace(Sx+1)
The exp(:) and log(-) operators above refer to matrix expo- nential and matrix logarithm, respectively. As S₃ is always positive semidefinite, the update in (13) incurs the computa- tional cost (n in the worst case. The general procedure for solving (SDR) is summarized Algorithm 1.
We first discuss how to construct supgradient estimators of F. By Danskin's 's theorem [7],
(F(S) = Conv Σ (S) **(S) E {.)s)
4.2. Theoretical Analysis
where (5) denotes the set of optimal solutions to the following OT problem:
In this subsection, we establish the complexity and perfor- mance guarantees for solving (SDR). Since the constraint set S2n is compact and the objective in (SDR) is continuous, an optimal solution, denoted by S*, is guaranteed to exist with finite optimal value. A feasible solution S E S2n,, is said to be 8-optimal if it satisfies the condition F(S) - ((S))F((S) < 8. Define the constant C' = max.,j ||,
min WET,
(M
(11)
The main challenge of constructing a supgradient estimator is to compute an optimal solution x² (S) E I (S). Since computing an exactly optimal solution is too expensive, we derive its near-optimal estimator, denoted as 8, and practically use the following supgradient estimator:
Theorem 4.4 (Complexity Bound). Fix the precision 8 >>> O and specify hyper-parameters
T 1602 log (m)) log((n)
log(2n)
vSS) Σ THe Mig 1.
(12)
c=
7
8²
4'
CVT
We adopt the stochastic gradient-based algorithm with Katyusha momentum in [82] to compute a e-optimal so- lution मे to (11). It achieves the state-of-the-an com- plexity O (n² See the detailed algorithm in Ap- pendix Next, we describe the main algorithm for soly- ing (SDR). Define the (negative) von Neumann entropy h(S) Exe(gn] (s( log (S), where {}(S)) are the eigenvalues of S. and define the von Neumann Bregman divergence
Then, the complexity of Algorithm for finding 8-optimal solution to (SDR) is
Next, we amalyze the quality of the solution to (SDR). Recall the exact reformulation (9) requires that the optimal solution to be rank-1 1 while the tractable relaxation (SDR) does not enforce such constraint. Therefore, it is of interest to provide theoretical guarantees the low-rank solution of (SDR), i.e., we aim to find the smallest integer k 2 1 such that there exists an optimal solution to (SDR) that is at most rank-k. The integer & is called a rank bound on (SDR), which characterized the following theorem.
V(S,, S2) = h(S)) h(S2) (S) - S2, T)) Trace(Si log S₁ Si log S2).
Iteratively, update S by performing mirror ascent with constant stepsize > 0:
Theorem 4.5 (Rank Bound on (SDR)). There exists an optimal solution to (SDR) of rank at most k. _ +
Skt1 = arg max ((u(Ss), + V(S.Sa) SeS..
6

---

Guarantees of KMS Wasserstein Distances
Problam (14)
Supp Find opiial (f.g) with Masilgg conmmint fi+ Boco S (M if [x]
See the general diagram that outlines this algorithm in Fig. 2 and the detailed description in Appendix J. We also provide its complexity analysis in the following theorem, though in numerical study the complexity is considerably smaller than the theoretical one.
Step 2: Greedy mok reduction algorihm
Step 2-2. Deternine if direction Q50" exists that Trace()) = and Mooyy Y =
Step 2-1: Figersassenpoiioonn of QaQ with digg(t
Yes
Uplae S (1)
Theorem 4.6. The rank reduction algorithm in Fig. sat- isfies that (T) for a 8-optimal solution to (SDR), it outputs another 8-optimal solution with rank at most k; (II) its worst-case complexity is ((n)).
Step 2-3-2: Find muximmat stepsize 5 such
that S (*)) - +.8Y satisfies (F) > 0, eg.
Step 2:3-1: Terminate and return 5
Aua (5*))) -0)
Figure 2. Diagram of the rank reduction algorithm. Here (-) de- notes the permutation operator on [n)). Step can be implemented using the Hungarian algorithm [41]. and Step 2-2 finds direction that lies in the null space of the constraint of Problem (14).
Additionally, by adopting the proof from Luo et al. [50], we show the optimality gap guarantee in Theorem 4.7. Al- though the approximation ratio seems overly conservative, we find (SDR) has good numerical performance.
Theorem 4.7 (Relaxation Gap of (SDR)). Denote by 4 (0.33)² an universal constant. Then
IV 2n 12++3 +3 3 } As a result,
en (SDR) I Optval(KMS) - OSDv).
max Optval(9)= S Trace (S) rank =1
F(S)
<Optval(SDR) S
max
F(S).
5. Numerical Study
SESt.
This section presents experiment results for KMS 2- Wasserstein distance that is solved using SDR with first-order algorithm and rank reduction (denoted as SDR-Efficient) Baseline approaches include the block coordinate descent (BCD) algorithm [75], which finds sta- tionary points of KMS 2-Wasserstein, and using interior point method (IPM) by off-the-shelf solver cvxpy [19] for solving SDR relaxation (denoted as SDR-IPM). Each instance is allocated a maximum time budget of one hour. All experiments were conducted on a MacBook Pro with an Intel Core 19 2.4GHz and 16GB memory. Unless other- wise stated, error bars are reproduced using 20 independent trials. Throughout the experiments, as Gaussian, with the bandwidth being the median of pair- wise distances between data points. Other details and extra numerical studies can be found in Appendices K and L..
The trivial rank bound on (SDR) should be 2n, as the matrix is of size 2n x 2n. Theorem 4.5 provides a novel rank bound that is significantly smaller than the trivial one.
Proof Sketch of Theorem 4.5. We first reformulate (SDR) by taking the dual of the inner OT problem:
{ { n (f
Σ
g)): 9 S (M,, S) Vi.j
max SESan f.gER"
(14)
By Birkhoff's theorem [8] and complementary slackness of OT, there exists an optimal solution of (f: g) such that at most n constraints of (14) are binding, and with such an optimal choice, one can adopt the convex geometry analy- sis from [43, 44] to derive the desired rank bound for any feasible extreme point of variable S. As the set of optimal solutions of (SDR) has non-empty intersection with the set of feasible extreme points, the desired result holds.
specify the kernel
Computational Time and Solution Quality. We first compare our approach to baseline methods in terms of running time and solution quality. The quality of given nonlinear projector is assessed by projecting test- ing data points from two groups and calculating their 2- Wasserstein distance. The experimental results, shown in the top of Fig. 3, indicate that for small sample size in- stances, SDR- IPM requires significantly more time than the other two approaches. Additionally, for small sample size instances, the running time of BCD is slightly shorter than that of SDR-Efficient stances, BCD outperforms running time. This observation aligns with our theoret- ical analysis, which shows that the complexity of SDR is Ó max(n, 8-))). lower than the complexity of BCD [75], which O( ((2883)) The plots the bottom of
It is noteworthy that Algorithm 1 only finds a nptr-optimal solution Sur of (SDR), which is not guaranteed to satisfy the rank bound in Theorem 4.5. To fill the gap, we develop a rank-reduction algorithm that further converts S1.r to the feasible solution that maintains the desired rank bound.
"Our motivation for rank-reduction is that it influences the quality of approximation the original non-convex optimization problem (KMS). Recall that globally solving (KMS) involves rank-1 constraint, and (SDR) relaxes it. When the solution to (SDR) is of rank > 1, we construct a feasible rank-1 solution by taking the leading eigenvector. However, if the solution to (SDR) is already low-rank, the rounding procedure typically yields feasible solution that is closer to the ground-truth rank-1 optimum.
However, for larger in-
in terms of
7

---

Guarantees of KMS Wasserstein Distances
Totbak Beam dit Car
Dateet Ovenn
OIPARE Datakk Atrlon Arkoaaee
3CD
BGD 5004NN
300-175
508-455
Eog
CIFAROT Darmet: Atzplase venoo Astomn/ako
0.20
3000
30403
2011 E.Frinet
Da Racoaedd 1QAN BEL.EAATTE Elrrann
SDR Eifiiwe
$2.5
2.8-
0.10
87.0
16
0.0
$3.0
310
200 Sumele
3.00
569
100
403
Sample
Sanyln
Figure 3. Comparison of SDR-Efficient with the baseline methods SDR-IPM and BCD in terms of computational time and solution quality. The columns, from left right, correspond to the synthetic Gaussian dataset (100-dimensiona)), MNIST, CIFAR-10. The top plots display the computational time, where the y-axis is labeled as "unbounded" if the running time exceeds the I 1-hour time limit. The bottom plots present the estimated KMS Wasserstein distance for each method.
Gaussian Covariance Shift
Gaussian Mixture Shift
100%
100%
Method
Testing Power
Testing Power
75%.
78%
Sinknom Div
80%
SWV
GBWW
0%l
0%
ME
0.01
0.
0.00
0.04
0.
0.b6
100
200
300
400
50o
600
700
snter-paramteer
Sample Size n
MNIST with distribution abundance change
CIFAR- 10 with distribution abundance change
100%
Method
Testing Power
Testing Power
KMS
75%
75%
Sinknom Div
SW
25%
25%
0%
ME
100
125
100
200
250
300
350
400
600
800
700
80o
800
1000
Sample Size
Sample Size
Figure Testing power with controlled type-1 error of 0.05 across four datasets. Figures from left right correspond (a) Gaussian covariance shift, (b) Gaussian mixture distribution shift, (c) MNIST, and (d) CIFAR-10 with distribution abundance changes.
Fig. 3 show the quality of methods SDR-Effi and BCD. We find the performance of solving SDR outperforms BCD, as indicated by its larger means and smaller variations. One possible explanation that BCD is designed to find a local optimum solution for the original non-convex problem, making it highly sensitive to the initial guess and potentially less effective in achieving optimal performance.
sample testing with other statistical divergences, such as (i) Sinkhorn divergence (S (Sinkhorn Div) [27], (ii) max- imum mean discrepancy with Gaussian kernel and me- dian bandwidth heuristic (MMD) [28], (iii) sliced Wasser- stein distance (SW) [10]. (iv) generalized-slice
Wasser- stein distance (GSW) [38], (v) max-sliced Wasserstein distance (MS) [16), and (vi) optimized mean-embedding test (ME) [36]. The baselines KMS, MS, ME partition data into training and testing sets, learn parameters from the training set, and evaluate the testing power on the testing set. Other baselines utilize both sets for evaluation. Synthetic datasets include the high-dimensional Gaussian distributions
High- dinenniiensii Hypothesis Testing, Then vali- date the performance of KMS 2- 2-Wasserstein distance for high-dimensional two-sample testing using both synthetic and real datasets. Baseline approaches include the two-
8

---

Guarantees of KMS Wasserstein Distances
KMS
SinkbormDi
MMD
SW
GSW
MS
ME
1060
1000
0.22
575
581
588
614
0..5
2300
2330
Figure 5. Top Left: Illustration of sequential data before and after the change-point. Remaining: Testing statistics computed from and
baseline approaches.
with covariance shift, or Gaussian mixture distributions. Real datasets include the MNIST [15] and CIFAR-10 [40] with changes in distribution abundance. The type-I error is controlled within 0.05 for all methods.
mance of various statistical divergences in generative model- ing by following the experiment setup in [39]. We compute the Fréchet Inception Distance (FID) scores generated fake images using convolutional neural network trained real images achieve satisfactory classification accu- racy. Lower FID score suggests better generative modeling performance. Experiment results reported in Figure 6, from which we can see that KMS Wasserstein distance can learn high-dimensional distributions using data.
We report the testing power for all these approaches and datasets in Fig. 4. For the Gaussian covariance shift scenario the MS method achieves the best performance, which can be explained by the fact that linear mapping is optimal to separate the high-dimensional Gaussian distributions. For the other three scenarios, our approach has the superior performance compared with those baselines.
1 6 à 3 0 4 3 0 9 P 9 8
6
6 /
7 0 0 0 9 /
6 9
l
0
9
0
f
7
0
0
Table Detection delay of various methods with controlled false
0
4
5 5
0
2
6
7 73332
0.05. Mean and standard deviation (std) are calculated based on data from different users.
alarm rate
3 7 6 7
2
9 5
9 9
2 0
0
2 3 5
3
5 9 q 9 9 7 6 7 5 6 6
8 7
9
9
Method
KMS
Sinkbor Div
MMD
sw
GSW
MS
ME
8 1
9 8 Q
OG
0
1 0
O
7
Mean
11.4
16.5
50.6
17.2
12.9
17.8
65.4
Std
5.56
4.4
39.5
8.7
6.4
9.2
25.7
2
6
6
60 9 3
6 7 8
6 4 7
9
2
/
0
8
0
/ 6 0
3
9
0
7
Human Activity Detection. We evaluate the performance of the KMS Wasserstein distance in detecting human activ- ity transitions quickly possible using MSRC-12 Kinect gesture dataset [21]. After preprocessing, the dataset con- sists of 10 users, each with 80 attributes, performing the action throwing/lifting before/after the change-pone at time index 600. We employ sliding window approach [81] with false alarm rate of 0.01 construct the test statistic each time index, which increases significantly when change- point is detected. Figure 5 illustrates the test statistics gener- ated by our method compared to baseline approaches. The experimental results, summarized in Table 1, show that our method achieves superior performance.
8
0
1
0
Figure 6. Plots from left to right correspond fake images ob- tained by generative models using (a) KMS Wasserstein, (b) Sinkhorn divergence, (c) MMD, (d) Sliced Wasserstein. (c) Gen- eralized Sliced Wasserstein, and (f) Max-Sliced Wasserstein. The FID scores 105.98, 1 114.36. 113.08. 113.92, 128.07, and 115.21, respectivoly.
6. Concluding Remarks
In this paper, we provided statistical and computational guarantees the KMS Wasserstein distance. Our numerical study validated our theoretical results and highlighted the exceptional performance of the KMS Wasserstein distance.
Generative Modeling. Finally, we examine the perfor-

---

Guarantees of KMS Wasserstein Distances
Acknowledgments
on pattern analysis and machine intelligence 39(9):1853- 1865.
This work is partially supported by DMS-2134037.
[15] Deng L (2012) The MNIST database of handwritten digit im- ages for machine learning research. IEEF signal processing magazine 29(6):141-142.
Impact Statement
[16] Deshpande I. Hu YT. Sun R. Pyrros A, Siddiqui N. Koyejo
This is a theoretical work on statistical and computational guarantees on the KMS Wasserstein distance. One of its societal impacts is its application to non-parametric two- sample testing. In practice, researchers can use two-sample testing to evaluate the effectiveness of medical treatments, discover economic disparities, detect anomaly observations, and more. We do not foresee any negative societal impact of this work.
S. Zhao Z. Forsyth D, Schwing AG (2019) Max-sliced Wasserstein distance and its use for gans. Proceedings of the IEEE/CVF conference computer vision and pattern recognition, -0656.
[17] Deshpande I, Zhang Z, Schwing AG (2018) Generative mod- cling using the sliced Wasserstein distance. Proceedings of the IEEE conference computer vision and pattern recog- nition, 3483-3491.
[18] Deza MM, Laurent M. Weismantel R (1997) Geometry of cuts and metrics, volume 2 (Springer).
References
[19] Diamond S. Boyd S (2016) Cvxpy: A python-embedded modeling language for convex optimization. Journal of Ma- chine Learning Research 17(83):1-5. -5.
[1] Altschuler J. Weed J, Rigollet P (2017) Near-lincar time ap- proximation algorithms for optimal transport via sinkhorn iteration. Advances in Neural Information Processing Sys- tems, 1961-1971.
[20] Flamary R. Courty N. Gramfort A. Alaya MZ, Boisbunon
A, Chambon S, Chapel I., Corenflos A, Fatras K, Fournier N. Gautheron L,, Gayraud NT, Janati H. Rakotomamoniy
[2] Balagopalan A, Novikova J. Mcdermott MB, Nestor B, Nau- mann I, Ghassemi M (2020) Cross-language aphasia detec- tion using optimal transport domain adaptation. Machine Learning for Health Workshop, 202-219 (PMLR).
A, Redko I, Rolet A, Schutz A, Seguy V, Sutherland DJ,
Tavenard R. Tong A, Vayer T (2021) Pot: Python optimal transport. Journal of Machine Learning Research 22(78):1- html.
8, URL http:/ / jmlr. wtgg/apprr///e /vtperrs//////// wwl.w222222222 451.
[3] Bartl D. Mendelson S (2022) Structure preservation via the
W Wasserstein distance. arXiv preprint arXiv:2209.0700 .20000088..........
[21] Fothergill S. Mentis H. Kohli Nowozin (2012) Instructing people for training gestural interactive systems. Proceedings of the SIGCHI conference on human factors in computing systems, 1737-1746.
[4] Barvinok AI (1995) Problems of distance geometry and con-
vex properties of quadratic maps. Discrete & Computational Geometry 13:189-203
[5] Ben-Tal A, Nemirovski A (2021) Lectures on modern convex optimization 2020. SIAM, Philadelphia.
[22] Fournier N. Guillin A (2015) On the rate of convergence in Wasserstein distance of the empirical measure, Probability theory and related fields 162(3):707 738.88
[6] Berlinet A, Thomas- C (2011) Reproducing kernel Hilbert spaces in probability and statistics (Springer Science & Business Media).
[23] Fukumizu K. Gretton A. Lanckriet G, Schölkopf B. Sripe- rumbudur BK (2009) Kernel choice and classifiability for rkhs embeddings of probability distributions. Advances in neural information processing systems 22.
[7] Bertsekas DP (1997) Nonlinear programming. Journal of the Operational Research Society 4833333343334
[8] Birkhoff G (1946) Tres observaciones sobre el algebra lincal.
[24] Gao R. Kleywegt A (2023) Distributionally robust stochas- tic optimization with Wasserstein distance. Mathematics of Operations Research 48(2):603-655.
Univ. Nac. Tucuman, Ser. A 5:14.
[9] Boedihardjo M (2025) Sharp bounds for max-sliced Wasser-
stein distances. Foundations of Computational Mathematics
[25] Gao R. Xie L., Xie Y. Xu H (2018) Robust hypothesis test- ing using Wasserstein uncertainty sets. Advances in Neural Information Processing Systems 31.
[10] Bonneel N. Rabin J, Peyré G, Pfister H (2015) Sliced and radon Wasserstein barycenters of measures, Journal of Math- ematical Imaging and Vision 51:22-45.
[26] Genevay A. Chizat L. Bach F, Cuturi M. Peyré G (2019)
Sample complexity of Sinkhorn divergences. Proceedings of the Twenty-Secon International Conference on Artificial Intelligence and Statisties,, volume 89, 1574-1583.
[11] Carriere M, Cuturi M. Oudot S (2017) Sliced Wasserstein kernel for persistence diagrams. International conference on machine learning. 664 673 (PMLR).
[27] Genevay A, Peyré G. Cuturi M (2018) Learning generative models with Sinkhom divergences. International Conference on Artificial Intelligence and Statistics, 1608-1617 (PMIR).
[12] Courty N. Flamary R, Habrard A, Rakotomamoniy A (2017) Joint distribution optimal transportation for domain adap- tation. Advances in neural information processing systems
[28] Gretton A, Borgwardt KM, Rasch MJ, Schölkopf B, Smola A (2012) A kernel two-sample test. J. Mach. Learn. Res.
30.
13(null):233 -773,,777,,37 ISSN 1532-4435.
[13] Courty N. Flamary R, Tuia D (2014) Domain adaptation with regularized optimal transport. Machine Learning and Knowledge Discovery in Databoses: European Conference,
[29] Hu Y. Chen X. He N (2021) On the bias-variance-con trade- off of stochastic optimization. Advances in Neural Informa- tion Processing Systems 34:221 19-22131.
ECML PKDD 2014, Nancy, France, September 15-19, 2014. Proceedings, Part I 14, 274-289 (Springer).
[30] Hu Y, Wang J. Chen X, He N (2024) Multi-level monte-carlo gradient methods for stochastic optimization with biased oracles. arXiv preprint arxiv:2408. /008.
[14] Courty N. Flamary R, Tuia D, Rakotomamonjy A (2016) Optimal transport for domain adaptation. IEEE transactions
10

---

Guarantees of KMS Wasserstein Distances
[31] Hu Y. Wang J, Xie Y. Krause A, Kuhn D (2023) Contextual stochastic bilevel optimization. Advances in Neural Informa- rion Processing Systems 36.
[49] Luise G, Rudi A, Pontil M. Ciliberto C (2018) Differen- tial properties of sinkhom approximation for learning with wasserstein distance. Advances in Neural Information Pro- cessing Systems 31.
[32] Hu Y, Zhang S, Chen X, He N (2020) Biased stochastic first- order methods for conditional stochastic optimization and applications in meta learning. Advances in Neural Informa-
[50] Luo ZQQ Sidiropoulos ND, Tseng P. Zhang S (2007) Approxi- mation bounds for quadratic optimization with homageneous quadratic constraints. SIAM Journal on optimization 18(1):1-
tion Processing Systems 33:2759-2770..
[33] Huang M. Ma S. Lai L (2021) Projection robust Wasserstein barycenters. International Conference on Machine Learning, 4456-4465 (PMLR).
28.
[51] Mairal J, Vert JP (2018) Machine learning with kernels. mines paristech, paris, france. URI htttss:/mmeberrss www.:ewww....iiieneeeeeeeeeeeeee mwnes-parisseehh master2217/mustez201...
[34] Huang M. Ma S, Lai L (2021) A riemannian block coor- dinate descent method for computing the projection robust Wasserstein distance. Proceedings of the 38th International Conference on Machine Learning, 4446-4455.
[52]
Mika S. Schölkopf B. Smola A, Müller KR, Scholz M. Rätsch G (1998) Kernel pea and de- be-noising in feature spaces. Ad- vances in neural information processing systems 11.
[35] Jiang B. Liu YF (2024) A riemannian exponential aug- mented lagrangian method for computing the projection ro- bust Wasserstein distance. Advances in Neural Information Processing Systems 36.
[53] Muandet K. Fukumizu K. Sriperumbudur B. Schölkopf B. al. (2017) Kemel mean embedding of distributions: review and beyond. Foundations and Trends in Machine Learning 10(1-2):1-141
[36] Jitkrittum W. Szabó Z, Chwialkowski KP, Gretton A (2016) Interpretable distribution features with maximum testing power. Advances in Neural Information Processing Systems 29.
[54] Nadjahi K. Durmus A, Simsekli U, Badeau R (2019) Asymp- totic guarantees for learning generative models with the sliced-wassersee distance. Advances in Neural Information Processing Systems 32.
[37] Kirchler M, Khorasani S, Kloft M. Lippert C (2020) Two- sample testing using deep learning. International Conference
on Artificial Intelligence and Statistics, 1387-1398 (PMLR).
[55] Nemirovski A, Juditsky A, Lan G, Shapiro A (2009) Robust stochastic approximation approach to stochastic program- ming, SIAM Journal on optimization 19(4):1574-160)
[38] Kolouri S. Nadjahi K. Simsekli U. Badeau R. Rohde G (2019) Generalized sliced Wasserstein distances. Advances in neural information processing systems 32.
[56] Nguyen K, Ho N (2023) Energy-based sliced Wasserstein dis- tance. Advances in Neural Information Processing Systems 36.
[39] Kolouri S. Pope PE, Martin CE, Rohde GK (2018) Sliced Wasserstein auto-encoders. International Conference on Learning Representations.
[57] Nietert S. Goldfeld Z. Sadhu R, Kato K (2022) Statistical, robustness, and computational guarantees for sliced Wasser- stein distances. Advances in Neural Information Processing Systems 35:28 179 28193.
[40] Krizhevsky A (2009) Learning multiple layers of features from tiny images. Technical Report TR-2009, University of Toronto, Toronto, Canada.
[41] Kuhn HW (1955) The hungarian method for the assignment problem. Naval research logistics quarterly 2(1-2):33-97.
[58] Pataki G (1998) On the rank of extreme matrices in semidef- inite programs and the multiplicity of optimal cigenvalues. Mathematics of operations research 23(2):339-3))
[42] Ledoux M (2006) Concentration of measure and logarithmic sobolev inequalities. Seminaire de probabilites XXXIII. 120- 216 (Springer).
[59] Patrini G, Van den Berg R, Forre P. Carioni M. Bhargav S, Welling M, Genewein T, Nielsen F (2020) Sinkhorn au- tocncoders. Uncertainty in Artificial Intelligence, 733-743 (PMLR).
[43] Li Y, Xie W (2022) On the exactness of dantzig- wolfe re- laxation for rank constrained optimization problems. arXiv preprint arXiv:22 1.0.1199.
[60] Paty FP, Cuturi M (2019) Subspace robust Wasserstein dis- tances. International conference on machine learning, 5072- 5081 (PMLR).
[44] Li Y. Xie W (2024) On the partial convexification of the low-rank spectral optimization: Rank bounds and algorithms. International Conference on Integer Progranming and Com- binatorial Optimization, 265 227 (Springer).
[61] Petzka H, Fischer A, Lukovnikov D (2018) On the regu- larization of Wasserstein gans. International Conference on Learning Representations.
[45] Lin Fan C. Ho N, Cuturi M. Jordan M (2020) Projection robust Wasserstein distance and riemannian optimization.
[62] Peyre G. Cuturi M (2019) Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning 1 (15-6))355-660
Advances in neural information processing systems 33:9383-
9397.
[46] Lin T, Fan C, Ho N. Cuturi M, Jordan M (2020) Projection robust Wasserstein distance and riemannian optimization. Ad- vances in Neural Information P Processing Systems, volume 33, 9383-9397.
[63] Ramdas A, Garcia Trillos N. Cuturi M (2017) On Wasserstein two-sample testing and related families of nonparametric tests. Entropy 19(2):47.
[64] Ramdas A. Reddi SJ. Póezos B, Singh A, Wasserman L (2015) On the decreasing power of kernel and distance based nonparametric hypothesis tests in high dimensions. Proceed- ings of the AAAI Conference on Artificial Intelligence, vol- ume 29.
[47] Lin T. Zheng Z. Chen E. Cuturi M. Jordan MI (2021) On projection robust optimal transport: Sample complexity and model misspecification. International Conference on Artifi- cial Intelligence and Statistics, 262-270 (PMLR).
[48] Liu F, Xu W, Lu J, Zhang G, Gretton A, Sutherland DJ (2020) Learning deep kernels for naramaraet two-sample tests. International Conference on Machine Learning, 6316-6326.
[65] Samadi S. Tantipongpipat U. Morgenstemn JH, Singh M. Vem- pala S (2018) The price of fair pea: One extra dimension. Advances in neural information processing systems 31.

---

Guarantees of KMS Wasserstein Distances
[66] Schrab Kim I, Albert M, Laurent B, Guedj B, Gretton
A (2023) MMD aggregated two-sample test. Journal of Ma- chine Learning Research 24(194):1-81
[67] Schrab A, Kim Guedj B, Gretton A (2022) Efficient aggre- gated kernel tests using incomplete u-statistics. Advances in Neural Information Processing Systems 35:18793-18807
[68]
Sidiropoulos ND, Davidson TN, Luo ZQ (2006) Transmit beamforming for physical-layy multicasting. IEEE transac- tions on signal processing 5446)22239-2511.
[69]
Sutherland DJ. Tung HY, Strathmann H. De S, Ramdas A, Smola A, Gretton A (2017) Generative models and model criticism via optimized maximum mean discrepancy. Inter- national Conference on Learning Representations.
[70] Wainwright MJ (2019) High-dimenssoonl statistics: A non- asymptotic viewpoint, volume 48 (Cambridge university press).
[71] Wang J, Chen M, Zhao T, Liao W. Xie Y (2023) A manifold two-sample test study: integral probability metric with neural networks. Information and Inference: A Journal of the IMA 12(3):1867-1888
[72] Wang J, Dey SS, Xie Y (2023) Variable selection for kernel two-sample tests. arXiv preprinl MxXx::2302.07115.
[73] Wang Gao R. Xie Y (2021) Sinkhorn distributionally robust optimization. arXiv preprint arXiv:2109. 11926
[74] Wang Gao Xic Y (2021) Two-sample test using pro- jected Wasserstein distance. 202/ IEEE International Sym- posium on Information Theory (ISIT), 3320-3325 (IEEE).
[75] Wang J, Gao R, Xie Y (2022) Two-sample test with ker- nel projected Wasserstein distance. Proceedings of The 25th International Conference on Artificial Intelligence and Statis- tics, volume 151.
[76] Wang J, Gao R, Xie (2024) Non-convex robust hypothesis testing using Sinkhorn uncertainty sets. 2024 IEFE Interna- tional Symposium
Information Theory (ISIT), 843-848
(IEEE).
[77] Wang J, Moore R, Xic Y, Kamaleswaran R (2022) Improving scpsis prediction model generalization with optimal transport. Machine Learning for Health, 474-488 (PMLR).
[78] Wang J, Xie (2022) A data-driven approach robust hypothesis testing using Sinkhorn uncertainty sets. 2022 IEEE International Symposium on Information Theory (ISIT),
3315-3320 (IEEE).
[79] Wendland H (2004) Scattered data approximation, volume 17 (Cambridge university press).
[80] Xie L. Xie Y (2021) Sequential change detection by optimal weighted & divergence. IEEE Journal on Selected Areas in Information Theory 1-1.
[81] Xie L. Zou S. Xie Y. Veeravalli VV (2021) Sequential (quickest) change detection: Classical results and new di- rections. IEEE Journal on Selected Areas in Information
Theory 2(2):494-514.
[82] Xie Y. Luo Y, Huo X (2022) An accelerated stochastic al- gorithm for solving the optimal transport problem. arXiv preprint arXiv:2203.0083
[83] Zhang Z, Wang M, Nehorai A (2019) Optimal transport in reproducing kernel hilbert spaces: Theory and applications. IEEE transactions on pattern analysis and machine intelli- gence 4277):1741-1754
12

---

Guarantees of KMS Wasserstein Distances
A. Comparison with Optimal Transport Divergences
Table 2. Time and sample complexity of empirical OT estimators in terms of the number of samples n. Here Ta, B represent two empirical distributions based i..d. n sample points R p denotes order of the metric defined the cost function standard Wasserstein distance.
Estimator
Complexity
Sample Complexity
[22]
Wp(an, P..)
Wasserstein distance
(n log m)
O(n (11d)/))
[26]
Spic (rre Pn)
Sinkhorn Divergence
O(n²) per iteration
O (12(1 el/3))
Max Sliced Wasserstein distance
[47)
MSs, (Pins P..)
((n d))
Ó 1/(max((k22)))
with k-dimennional projector
[57]
SWp(0..P.. P..))
Sliced Wasserstein distance
((dd)
Ó(n 11((2))
[75] and this work
KPM Vin)
Kernel Max Sliced Wasserstein Distunce
((n2d")
O(n (1((2))))
In Table 2, we summarize the time and sample complexity of various OT divergences. Notably, the sample complexity of the standard Wasserstein distance suffers from the curse of dimensionality. While the Sinkhorn divergence mitigates this issue, its sample complexity depends on ed/2, which can be prohibitively large when the regularization parameter is very small. For Max-Sliced (MS) and Sliced Wasserstein distances, their sample complexity is independent the data dimension. However, they rely on linear projections to analyze data samples, which may not be optimal, particularly when the data exhibits nonlinear low-dimensional structure. This limitation motivates the study of the KMS Wasserstein distance by Wang al. [75] and in our work. In Example 2.5 and Section numerically validate this observation and demonstrate its practical advantages.
B. Proof of Theorem 3.2 and Corollary 3.3
The proof in this part relies on the following technical results.
Theorem B.1. (Finite-Sanple Guarantee on MS 1-Wasserstein Distance on Hilbert Space, Adopted from (9, Corollary 2.8]) Let 5 E (0, 1), and be a probability measure on a separable Hilbert space H with Jn | du(x) < co.
Let XI...., X, be ii.dd random elements of sampled according to μ,, and pn = Σi= 5x,, then it holds that
EMS; (n, SC ( ||0||2+26 [10112+2ddpxxz) 1/(2425)
(66) 1/2
where C > is a universal constant. Theorem B.2 (Functional Hoefffding Theorem [70, Theorem 3.26)). Let F be class of functions, each of the form h: B + R, and XI....,X. Xn be samples ii.d. drawn from μ on B. For i E [n), assume there are real numbers di S bi such that h(x)= E [as,h, for any x E B, h E FU F Define
1 Σ
L
sup hEFU(-F)
(b., - a2,1)
n i=1
For all 0, holds that
(_
.182
P
sup NEF
h(X))
2 E
sup he
h(X))
+ 8 < exp
4L2
We first show the one-sample guarantees for KMS y-aasserstein distance.
Proposition B.3. Fix p E [1, oc), error probability (0, 1), and suppose Assumption 3.1 holds. Let C. 2 1 be universal constant. Then, we have the following results:
RKMS (mn) ) A(2C1/9).n-1)) -1(1(22))
13

---

Guarantees of KMS Wasserstein Distances
(II) With probability least 1 a, holds that
1/8
Pm) $21-1/P C+ (C+4 log _ a
n 1/(2p)
Proof of Proposition B.3. Recall from (3) that
k. -M,,, (0#11, $*U)
Therefore, it suffices to derive one-sample guarantees for MS, $ $ )
(I) Observe that under Assumption 3.1, we have
A z ((x = (Ka, K₂) = 11 1
and therefore |0(x))1
A. V2 B. In other words, for every probability measure μ B, the probability measure # is supported on the ball in H. centered the origin with radius A. By Theorem B.1 with = 1, we obtain
[| I <
AC
IKMSS (a)) = IEMS (4 & #
VTi
Since @ a a and P are supported on the ball of H. centered at the origin with radius A, holds that
MSp (05p ( #Hmn) s [MS. (### f+ (2A))-1)
1/p
In other words,
S (M, (M, (2A))-1) 1/p
(15)
It follows that
EKMSS,(H,PP) EMS, @#H,, 2
(2A))-1111 1/0
s E MS.
@#11 1m
( p+ (2A))-1) [13 1/p
1 MS
{ vm (2A)a-1)
AC
} 1/8
21 AC /v . 1/(pp)
(II) For the second part, we re-write KCMS (M, (in) with in 1 n
Σn 8r, using the Kantorovich dual reformulation of
OT:
Σ (ols(x)) -Famm
sup
JEM: lf
g is L-pisschizz with (0) =0
where the additional constraint g(0) = 0 does not impact the optimal value of the OT problem. In other words, one can represent
(MM. sup hEs
M M(rs)
where the function class
5 x +((()))) (a : g 1-Lipschitz with g(0) =0, ʃE H, 115 $1 Consequently, for any x,
lo(f()))) lo((())) -(0) < 1f(x)) KV. (15la Kall S A.
14

---

Guarantees of KMS Wasserstein Distances
One apply Theorem B.2 with F S. diS -A Eow [o([(x))), bud h(x)) (((())) E [l([((()))))) to obtain
A Ee>plo((())) where
n262
no2²
((I.MSS(AR (m)) 2 (I +0] - exp
4(2A)²
= exp
16A42
Or equivalently, the following relation holds with probability at least
1 a (M, - E [KMS + 4An 1/2 V log I An2
1
C+
log
a
By the relation (15), we find that with probability at least a,
K.MSp(1 Fmm)
1/p
An 1n22
C log
1
(2A))-1) 1/p = 21-1//A
(C+4/ log
11(2p)
n
We now show the proof of Theorem 3.2(II). By the triangle inequality, with probability least 2o, it holds that KMM Fn) & K.MS,((,T.)) + K.M Dn) + x.MSs(((,))
1/v
22.21-1/DA
C C+4y1ogg log
-11((22)) + k.M
1/8
S 4A
C+ 4 4 log
n -1((2) + K.M,S(u.))
Then, substituting with a/2 gives the desired result.
Proof of Corollary 3.3. It remains to show the type-11 II risk when proving this corollary. In particular,
Type-II Risk
PH, (KMS Pn) A(n,a))
Pn,
K.MS (n) K.M - ((n,a))
I PH, -)
- KKms Bwn) 2 K.MSp((,u) - A(n,a))
-KMM) K.MSp(u,, (nn
where the last relation is based on the Markov inequality and the assumption that K.S v) - A(n, ar) > 0. Based on the triangular inequality, we can see that
((K.MS)(H,,) (Pn, Pw)) (H, EEM.M 2ACH/P n
Combining these two upper bounds, we obtain the desired result.
C. Sufficient Condition for Positive Definiteness of Matrix G
To implement our computational algorithm, one needs to assume the gram matrix
= [K("" x"). K(x",
((z"."0 yy") KKy" x²) K ((y",')) y)"]
to be strictly positive definite. By the Lemma the Schur complement (see, e.g., [5, Lemma 4.2.1)), It can be showed that its necessary and sufficient condition should be
G' = [K("", x²), K(x", "<=","")) ; ((y" K(y","")
15

---

Guarantees of KMS Wasserstein Distances
strictly positive definite. By Wendland [79], this requires our data points ("1:-.., x", . Un are pairwise distinct and K (x,,, y) is of the form ( ((x y); with ((-) being continuous, bounded, and its Fourier transform is non-negative
and non-vanishing. For instance, Gaussian kernel K (as, y) = -1z--vl//0" or Bessel kernel Kxx, y) (c + || - 2 Rd 3 d/2 satisfies our requirement.
D. Reformulation for 2-KMS Wasserstein Distance in (KMS)
In this section, we derive the reformulation for computing 2-KMS Wasserstein distance:
max FEH. IISR
min meT (55/M
Σ
TTylf(z)-
(16)
Based the expression of f in (7), we reformulate the problem above
max aR,
min T
Σ
mij
Σ LE(m]
ax
S (2is, xii)
E
a (ysr 31)
(17a)
1E(m)
subject to the constraint
2 a (.,3t)
Σ ag K( (,, - Σ
1=1
n
d2 (2,K(;;,;)) Σ
ay.iK(., (.,88e) Σ
aniiK( (e,,i((,,))) (,,z;) Σ ag (.
1=1
E
an (K(,)))) K (x,)))))
+
Σ dy ( (k(+ K ((,,s))) - Σ
a2 o (K(; K(s S..
(JE(M]
4.fE/m]
(17b)
One can re-write (17) in compact matrix form. we define
[az; Ayl:
Mr = [K(*i, 01) - (yi, x)) (i))))))) (K(y,, 31) K K(r,, 91 )[]]]
G [K(",, xx"), - K("", "a"",yl yy"); x)", K KV"", w²)] E I8 2n
Problem (17) can be reformualted as
max FERAn
min meT (JE(M]
Σ The S5 M
TTGs GG $1
(18)
Take Cholesky decomposition G = UTT and use the change of variable approach take U-1s, Problem (18) can
be further reformulated
max W
min FCT
Σ Tes ((, Ww 1).
(19)
1.jE(n]
After defining Mrj U Mi and observing that the inequality constraint Tw will become tight, we obtain the desired reformulation in (9).
E. Proof of Theorem 4.2
The general procedure of NP-hardness proof is illustrated the following diagram: Problem (9) contains the (Fair PCA with rank-1 data) as a special case, whereas this special problem further contains (Partition) (which is known to be NP-complete) as special case. After building these two reductions, we finish the proof of Theorem 4.2.
16

---

Guarantees of KMS Wasserstein Distances
Give integers di. AN determine whether binary variables {er (-1,,, exist such that Σ (422
min w A,A, we |=] [n(n]
min Σ Fio (Al
DIX Iulb-1
(9)
(Fair-PCA with rank-1 data
(Partition
Figure Proof outline of Theorem 4.2
Claim 1. Problem (9) contains Problem (Fair PCA with rank-1 data).
Proof of Claim 1. Given vectors A1.... An, we specify
M1 & {M1,1 M1....... M1,, (AI....,AA)).
and M. & {M11, Mr....... .., M. For instance, M2.: (An, AI,..., An-1) }. For the inner OT problem in (9), it suffices to consider deterministic optimal transport iT, i.c.,
M,, }, 2, ... is specified by circularly shifting the elements in M1,; by i positions.
1/n, if if)= 0,
Tijj
otherwise
for some bijection mapping a [n]] + [n]]. The cost matrix for the inner OT is actually a circulant matrix:
(Azw)2 (Anw)² (Aww)²
(Aww)²
(Anw)2
(An-1w)2
(MM(W))))
(A2w)² (Azw)²
(Aww)
When considering the feasible circularly shifting bijection mapping (e.g., ai) = (i + j) mod n, Vi E [n] for j 0, ,,1,...,22 1), we obtain the upper bound on the optimal value of the inner OT problem in (9):
min RET
Σ Tis (M. 4
min (A/w))2 min rEin
wT A'A,A,W w.
iE(n]
On the other hand, for any bijection mapping a, the objective of the inner OT problem in (9) can be written as convex combination of (2fw)2,.... ...,,, w?? and thus,
Σ [mn EE..
min
min OER a.
Σ [o:((AW))}]
(4Tw)) EE(n)
min
Since the upper and lower bounds match with each other, we obtain
Σ [MF min w A, Te,,
min
sE(n]
and consequently,
max w: ||0||==1
min =er,
Σ The (M (w)
max w: [
min w A, (E(()
1.j
which justifies Problem (9) contains Problem (Fair PCA with rank-1 data).
Claim 2. Problem (Fair PCA with rank-1 data) contains Problem (Partition).
noteworthy that Claim has previously been proved by [68]. For the sake of completeness, we provide the proof here.
17

---

Guarantees of KMS Wasserstein Distances
Proof of Claim 2. Consider the norm minimization problem
P min
1
min se|n]
(w A,) 1
(20)
and the scaled problem
max
min sE(m]
(wTA)) A.)² 1||12 }
(21)
We show that Problem (Fair PCA with rank-1 data) is equivalent to (21), whereas (21) is equivalent (20). Indeed,
For the first argument, for any optimal solution from Problem (Fair PCA with rank-1 data), denoted as w one can the scaling to consider a* v Pw,, which is also optimal to (21), and vise versa.
For the second argument, let w1, W2 be optimal solutions from (20), (21), respectively. Since P is the optimal value of (20), can check that w1 is feasible solution to (21). Since minig[n] (w A,) 2 1, by the optimality of w2, it holds that min] (W²'A.)* A.)² 2 1, i.e., w2 is feasible solution to (20). Since ||22| P. is an optimal solution to (20). Reversely, one can show wi is an optimal solution to (21): suppose on the contrary that there exists W1 such that || 11 = P and min;e(n] (w (TFA,) A,)² > min;e[n] (wT A,)² 2 1, then one can do a scaling of 11 such that minge[n] (of A = 1 whereas ||<1 112 > P. which contradicts to the optimality of P. Combining both directions, we obtain the equivalence argument.
Thus, it suffices to show (20) contains Problem (Partition). Define (ai)ke(n), Q = I, + aa and assume Q admits Cholesky factorization Q S. Then create the vector A, S-T eg, where e, is the i-th unit vector of length Then, it holds that
(20)
min [|
wi) 2 21}
iE(n
min
||50 [
min(x 2 }
iE|n]
T Qu: o 1
min
2
-min Σ
Σ
x 2 }
where the second equality is by the change of variable = S- the third equality is by the definitions of S and e,, and the last equality is by the definition of Q. The solution to Problem (Partition) exists if and only if the optimal value to Problem (*) equals n. Thus, we finish the proof of Claim 2.
F. Algorithm that Finds Near-optimal Solution to Optimal Transport
In this section, we present the algorithm that returns c-optimal solution to the following OT problem:
min SET.
Σ
(22)
1
where (cighis the given cost matrix. Define ||C||o max,, C.j- In other words, find ग E F, such that
optval(22)
m.jo S optval(22) +e. Σ
i.j
Entropy-Regularizee OT.
The key to the designed algorithm is to consider the entropy regularized OT problem
min r
Σ Eap + +n
maj (.g(
4
18

---

Guarantees of KMS Wasserstein Distances
whose dual problem is
n Σ 5m000)
min #ER
G(o)
(23)
where
h2(2) = 77 log
Σ exp
1
Ci.j -
1 Σ 2 + 711 + log
1
Given the dual variable e IR,, one can recover the primal variable using
1
exp
π (u))
Ej' exp
Algorithm 2 essentailly optimizes the dual formulation (23) with light computational speed.
Algorithm 2 Stochastic Gradient- based Algorithm with Katyusha Momentum for solving OT [82]
Input: Accuracy > 1) 8 8og
6' maxi, ci. maximum outer iteration Touts and maximum inner iteration
2: Take (yo, 20, Ao Ao. Co, Do) = (0, (0,0,0,0,0,0)
3: for 0,... ...., Toul do 2
4: Tiet 144: Tt. ST.
5: Ut V(((+t)
6: for j 0...., ,T' do
7:
==
8:
AR+1 T1,1 + 17, + ( 71,)
9:
Sample uniformly from [n]. and construct
Ha+1 ur + (vhs(Ae+1) - vhh(2e))
10:
Update Zx+11 2k 7u- H2+1/2 and Wk+1 = Akt1 1//14+1/9
11:
end for
12:
Update Ae+1 + -e= VeT+j
13:
Sample Ne uniformly from (Ae++1.... Mer+1....,Aerr Aer+T] ..rrr)))))))))) and take D. = Dr + vec(T(R)))/)
14:
C. = C2 1/T1,
T = D./C.
end for
17: Query Algorithm 3 to Round म Tr.. to @ such that #1, 11,, and 1 3 1,
18: Return R
Algorithm 3 Round to r, ([1, Algorithm 2])
Input: R B
2: X diag(E1,..., Rn), with x, min (1, mn(ry) }. Here (T)) denotes the i-th row sum of
3: n' Xx.
4: Y
with y, = min(1, medtrr) ). Here e, (T') denotes the j-th column sum n!.
5: T" T'Y.
6: e, = 11, - = 11, - c(TT), where
r(a") = (r:(1") c(x") ())())))))))
7: Return x" +,,e le/ll.
19

---

Guarantees of KMS Wasserstein Distances
Theorem F.1 (182, Theorem 3)). Suppose we specify Tou = oo ( ). T n, the number of total iterations (including outer and inner iterations) of Algorithm 2 is =0 v with per-iteration cost (n)) Therefore, the number of arithmetic operations of Algorithm 2 for finding e-optimal solution is O n=||C] -. vinn
G. Proof of Theorem 4.4
To analyze the complexity of Algorithm 1, we first derive the bias and computational cost the supgradient estimator v(S) in (12).
Lemma G.1 (Bias and Computational Cost). The following results hold: (T) (Bias) v(S) corresponds to the gradient of F(S) = Eng Re M...MMM.... where ß is defined in (12) and [(S) F(S) S c;
(IT) (Cost) The cost for computing (12) O (C n² Vlogne with (()) hiding some universal constant.
Next, we analyze the error of the inexact mirror ascent framework in Algorithm 1.
log(2n) Lemma G.2 (Error Analysis of Algorithm 1). When taking the stepsize 7 OVT
the output SuT from Algorithm
satisfies
0 0 F(S*) -(S))) 2e 201 log(2n)
Combining Lemmas G.I and G.2, we obtain the complexity for solving (SDR).
Proof of Lemma C.I. For the first part, it is noteworthy that v(S) is associated with the objective
F(8)= = Σ t.s MM.
4
where Ri is the c-optimal solution to By definition, it holds that
F(S)
min xef
Ti (M. Mig, E
S F(S)-F((F) S.
The second part follows from Theorem F.1.
The proof of Lemma G.2 replies on the following technical result. Lemma G.3 ([55)). Let (Sk) be the updating trajectory of mirror ascent aiming to solve the maximization of G(S) with S E S2n i.e.,
Sk+1 arg max (v((S)), + Ss), k .. 1. SES
Here v(S) is supgradient of G(S), and assume there exists M. 0 such that
[[v(S)) == Trace(v(SS) S M.. VS E Szn-
Let S1.7 Et= Sk. and 5* be maximizer of G(S). Define the diameter
D2
max h(S)
min h(S) SESmn
log 1og(2n).
SESan
For constant step size
D2 S.n log(2n)
7
M.VT M.VT
it holds that
4 Algg(n) (2))
0 0 G(S") G(STT) & Ms
T
20

---

Guarantees of KMS Wasserstein Distances
Proof of Lemma G.2. Let S and S* be maximizers of the objective F()) and F()) then have the following error decomposition:
F(S") F((1:T)
=IF(S)) - F(S*)] + [F(S)) - + (()r))))))))) -
52c [F(S*) - -F(S*] -F(5.))((3L.)))
22c + [F(St) - F(SI:r)),
where the first inequality is because || Floe and
[[F(S*) F(S*)] S S. (. - F(SIT)]] S
and the second inequality is because F(St)-F(S1n) S 0. It remains to bound [ (((i)))) It worth noting that
[|u(S)] = Σ
T || M, M S Σ Tte C ==
1
Therefore, the proof can be finished by querying Lemma G.3 with M, = C and stepsize log(2n) CVT
Proof of Theorem 4.4. The proof can be finished by taking
such that
log (2
2c S
2C
2
In other words, we take 6 and T 16C (2g)
We follow the proof of Lemma G.2 to take stepsize y
Ing(2n)
62
C/T
H. Proof of Theorems 4.5 and 4.7
We rely on the following two technical results when proving Theorem 4.5. Theorem H.1 (Birkhoff-von Neumann Theorem [8]). Consider the discrete OT problem
min Σ Tes
(24)
TE N 1
There exists an optimal solution that has exactly entry of 1/n in each and each column with all other entries 0. Theorem H.2 (Rank Bound, Adopted from [44, Theorem 2] and [43, Lemma 1]). Consider the domain set
D = S s Trace(S) = 1}
and the intersection of N linear inequalities:
E 8 { (S.A)) > b, E [N]].
Then, any feasible extreme point in D E has a rank at most [/22 + 9/4 -322.. Such a rank bound can be strengthened by replacing N by the number of binding constraints in E.
Proof of Theorem 4.5. By taking the dual of inner OT problem, we find (SDR) can be reformulated
1 - + 9:): f. S (M, 8), Vi.j E im)
max SES
in]
(25)
f.gER*
Let S* be the optimal solution variable the optimization problem above. Then for fixed S*, according to Theorem H.1 and complementary slackness of OT, there exists optimal solutions ("*,g") such that only constraints of n² constraints
21

---

Guarantees of KMS Wasserstein Distances
in (25) are binding. Moreover, an optimal solution to (SDR) can be obtained by finding feasible solution the following intersection of constraints:
Find S E S2n n & - {s: S: fi + S Mi
S), 8), i,j E Irl].
By Theorem H.2, any feasible extreme point from San n has rank at most 1 1+ 2n - 2/] feasible extreme point to satisfy the requirement of Theorem 4.5.
|/22
Thus, we pick such
Proof of Theorem 4.7. Recall that
(KMS)
max
Σ 72
+:): J+ S (Mi,
Vi,j E |u]
S2 0, Trice(S) (.n)=_.nn(()---) =1,
fIDER"
and
1
(SDR)
max Trace( S.gER"
f S (Mi Vi,j E [n] +f
n
It is easy to see Optval(KMS) - OSDR). On the other hand, let (S, be an optimal solution to (SDR) such that
(S) S k A 1+
+ 2n
+
1 -
Next, take
C
5 S),, ==
5 = EET
It can be seen that S 2 0, Trace(S) = 1, (Snk((S) = 1. Then, for any E (0, 1] and f >> 0, it holds that
Pr (Mr, M. 1. 5) 2 e. (fi +f+))). Vi,j E [ml]
2 Pr (Mis 1. 5) 8 8. (Ma, M 1. 5), vi,j E [ml]
=
-Pr
M.. 5. 5) 2 E( (1 Vi.e [
=
=pr
(Mis 1: (5 ||<|| 2 B(MM, MC Vi.j [m]
PP
(Mi,
" > μ
_. MK <("),, e.. 1|<112 2 11
21 Σ
Pr
(c é M(, <<>>]} Pr[1Kcl12 } P
(6.))lmn)
21 n²²
1 μ.
As long take f 33/ 100, 4 (33/100)3/n1 it holds that
Prf RT 3) e. -ft+ Vi.j E [ml] 2 1/100.
In other words, there exists & such that
(s fi) + (e. (e- <ET",, Vi,j E [m] ||E||2 1.
This indicates that (55 E f, 9) is a feasible solution to (KMS), and consequently.
OptV
(e fite 9.) = (KMS).
22

---

Guarantees of KMS Wasserstein Distances
Algorithm 4 Rank reduction algorithm for (SDR)
1: Run Algorithm I to obtain 6-optimal solution (SDR), denoted as S.
// Step 2: Find n binding constraints
2: Run Hungarian algorithm [41] to solve OT (11) with S = S, and obtain an optimal assignment with dual optimal solutoin gf*
[n] - [n] together
// step 3-9: Calibrate low-rank solution using greedy algorithm
3: Initialize 8* 1
4: while 8 > do
5: Perform eigendecomposition S QAQT, where A = diag(A1...... A,) with rank(
6: Find direction Y = QQQQT, where A E S" is some nonzero matrix satisfying
Trace A) 0, (Q (M =0, Vi E [m].
7: If such Y does not exist, then break the loop.
8: Take new solution S(0*)= = + 8 where
8 = arg max {5: Amin(A + + 2 0).
820
9: Update 3 = S(8*)
10: end while
11: Return
Rank Reduction Algorithm
In this section, develop rank reduction algorithm that, based the near-optimal solution (denoted as S) returned from Algorithm 1, finds an alternative solution of the same (or smaller) optimality gap while satisfying the desired rank bound in Theorem 4.5.
Step (i): Find n binding constraints. First, we fix = S (14) and finds the optimal solution (f",g") such that only n constraints out of n² constraints are binding. It suffices to apply the Huangarian algorithm [41] to solve the following balanced discrete OT problem
max f
Σ U++g:):
f Oj S Cij
= min wET,
Cig 1,=1
where q.j = (M,, S). The output of the Huangarian algorithm is a deterministic optimal transport that moves probability mass points from the left marginal distribution of ı to the right, which is denoted as a bijection that permutes [n] to [n]. Thus, these n binding constraints are denoted as
f + 906) MM 8), i [n].
We denote by the intersection of these n constraints as E,,.
Step (ii): Calibrate low-rank solution using a greedy algorithm. Second, let us assume 3 not an extreme point of S2n E,, since otherwise one can terminate the algorithm to output S following the proof of Theorem 4.5. We run the following greedy rank reduction procedure:
(I) We find direction Y # 0, along which á remains to be feasible, and the null space of S is non-decreasing.
(II) Then, we move I along the direction Y until its smallest non-zert eigenvalue vanishes. We update Š to be such a new boundary point.
23

---

Guarantees of KMS Wasserstein Distances
(III) We terminate the iteration until movement is available.
To achieve (I), denote the cigendecomppssiton of S with rank(S) as
A 0
$ (Q 0)
(QT
0) QAQ
0
where = diag(A1,..... Ar) with 2 2 d > 0 and Q Raaxr, To ensure S SY
S2a &n while
Null(S àY) 2 Nu)l(S), for some stepsize 8 > 0, it suffices to take
Y= (Q 0) A 0 0
(Q 0) QAQ""
where 8" {0} is symmetric matrix satisfying
Trace(A) (A) O, [M..M... 0, i [
To achieve (II), it suffices to solve the one-dimensional optimization
8 arg max {o: Amin (A + +4) 2 0). 820
(26)
where (-in()) denotes the smallest eigenvalue of a given matrix. the optimization above admits closed-form solution update. Let eigenvalues of A be X 2 ... 2 X,. It suffices solve
8 = arg max {5: min (A + 5xf) 2 0).
820
sefr]
As long as X, 2 O, we return 8 0. Otherwise, let an index such that X 20> Ma11 We take 8 maxicySr the desired optimal solution.
The overall algorithm is summarized in Algorithm 4. Its performance guarantee is summarized in Propositions and Theorem 4.6.
J. Proof of Theorem 4.6
The proof of this theorem is separated into parts.
Proposition J.1. The rank of iteration points in Algorithm strictly decreases. Thus, Algorithm 4 is guaranteed to terminate within 2n iterations.
Proof. Assume on the contrary that rank(S(ōo)) rank($) = r. Since 3(6*) = Q(A + A)QT. the positive eigenvalues of 5(8*) are those of the matrix A + 8 A. According to the solution structure of (26), this could happen only when A + 8 A > 0, i.e., either 8* = 0 or A > 0. For the first case, this algorithm terminates. For the second case, since Trace( A) = 0, A E S", it implies that A = O, which is a contradiction.
Thus, the rank of the iteration point reduces by at least in each iteration.
Proposition J.2. Let S be the output of Algorithm 4. Then, it holds that
(I) S* is a 8-optimal solution to (SDR).
(II) The rank of S" satisfies
9 4
3
rank(S*)) < +
2n +
2
24

---

Guarantees of KMS Wasserstein Distances
Proof. Recall the solution $ obtained from Step of Algorithm 4 satisfies
F(S)
min TE
Σ
Ti. (M,
1.
fit S (MI 3) (SD)
{ 5.
max f.geR
Since Step 2 of Algorithm 4 solves the OT problem exactly, obtain
Σ ([) F(S) (SDR) -
Since Step 3-7 always finds feasible solutions to the n binding constraints
+ Dos M. S), i i [n],
for any iteration points from Step 3-7, denoted as S, the pair (f. .g) ) guaranteed to be the d-optimal solution to (14), a reformulation of (SDR). Hence we finish the proof of Part (I).
For the second part, assume on the contrary that rank(S*) 2 1+ 2n - It implies + < r(r + 1)/2.
Recall that Step 6 of Algorithm 4 essentially solves linear system with n + constraints and r(r + 1)/2 variables, nonzero matrix A is guaranteed to exist. Thus, one pick sufficiently small 5 > 0 such that Amin (A 5A) 2 0, which contradicts to the termination condition 8 = 0. Thus, we finish the proof of Part (II).
Combining both parts, we start to prove Theorem 4.6.
Proof. Algorithm 4 satisfies the requirement of Theorem 4.6. For computational complexity, the computational cost of Step 2 of Algorithm 4 is (n²)) In each iteration from Step 3-7, the most computationally expansive part is to solve Step 6, which essentially solves a linear system with n - constraints and r(r + 1)/2 variables. The conservative bound r S 2n. Hence, the worst-case computational cost of Step (which can be achieved using Gaussian elimination) is
O(n 1 1+ 1)/2) (n + 1) O(n*).
Since Algorithm terminates within at most 2n iterations, the overall complexity of is O(z²).
K. Numerical Implementation Details
K.1. Setup for Computing KMS Wasserstein Distance
When implementing mirror ascent algorithm, for small sample size (n < 200), the exact algorithm adopted from https: /ttptooottoo.tt.ooooo github. io/ to solve the inner OT; whereas for large sample size, algorithm adopted from ht / /github. com/ Yi wi1iinxxe27/PDASSS to solve this subproblem. For the baseline BCD approach, we implement using the code from gi github.... www ttt.b//Radteeeaaaadadad httdddd/dXRXee_/eedee
the approximation
i..
K.2. Setup for High-dimeusional Hypothesis Testing
For baselines Sinkhorn Div, SW, MS, we implement them by calling the well-established package POT [20]. For the ME baseline, we implement it using the code from https: / /gttttb.... com/ wii ttttttaaat..
http://tcc///tttttttt///tttteeee w/interpretaale--teeeeeeee teet..
Next, we outline how to generate the datasets for two-sample testing experiments in Fig. 4:
The Gaussian covariance shift dataset was generated by taking
(0, Ia), NO pF).
where the dimension d = 200, the sample size n = 200, and E is the all-one matrix. We vary the p from 0.06.
25

---

Guarantees of KMS Wasserstein Distances
MII The Gaussian mixture dataset was generated by taking
1 1 - == N00, Ta) + N la, Ia), v = 2
1
la + - 2 1a la + 0.05Æ),
where the dimension d = 40, and we vary the sample size n from 100 to 700.
(III) For MNIST CIFARIO examples, we take j the uniform distribution subsampled from the target dataset, denoted
as Pdata- We take as the one having change of abundance, i.e., V 0.85paata + 0.15p of lahel where
Piata label corresponds to the distribution of a subset of the class with label 1.
Finally, we outline the procedure for practically using the KMS Wasserstein distance for two-sample testing, based on the train-test split method:
We first do the 50%-50% teining-testing data split such that x." Tr U To and y" = T Uy
(II) Then we compute the optimal nonlinear projector f for training data (xx*, "h). We specify the testing statistic as the Wasserstein distance between projected testing data "
(III) Then we do the permutation bootstrap strategy that shuffles (2: Tc) for many times, e.g., L = 500 times. For each time I, the permuted samples are (xr(e)> M(()). and we obtain the testing statistic the Wasserstein distance between projected samples (using the estimated projector f. denoted as W(f
Te (f) (+)
(IV) Finally, we obtain the threshold as (1 - quaantiee of the testing statisties for permuted samples, where the type-I error a 0.05.
K.3. Setup for Human Activity Detection
The MSRC-12 12 Kinect gesture dataset contains sequences of human body movements recorded by 20 sensors collected from 30 users performing 12 different gestures. We pre-process the dataset by extracting 10 different users such that in the first 600 timeframes, they are performing the throwing action, and the remaining 300 timeframes, they are performing the lifting action.
K.4. Setup for Generative Modeling
We follow Deshpande et al. [16] to design the optimization algorithm that solves the problem
min D D(Pdata: (fe) PPoisee)
where Piata denotes the empirical distribution of MNIST dataset, Paoise denotes the Gaussian noise, and (Jo) #roise represents the distribution of the fake image dataset. We specify fe as a 4-layer feed-forward neural-net with leaky relu activation, and denotes its weight parameters. We train the optimization algorithm in 30 epoches. From Fig. 6, we observe that the KMS Wasserstein distance provides fake images that are more closer to the ground truth compared with the Sliced Wasserstein distance.
L. Additional Numerical Study
L.1. Experiment on Significance Level
Here we add the experiment in the following table to show that as long as we take 0.05, the practical type-I error of KMS Wasserstein test is guaranteed to be controlled within 0.05.
Table 3. Type-I Error for two-sample testing with Gaussian mixture datasct.
Method
20
40
80
160
180
200
KMS
0.061 + 0.015
0.055 0.012
0.043 + 0.014
0.059 0.011
0.042 - 0.014
0.062 = 0.015
26

---

Guarantees of KMS Wasserstein Distances
L.2. Rank Reduction Algorithm
Recall that Theorem 4.5 provides the rank bound regarding some optimal solution from SDR. In this part, compare the rank of the matrix estimated from Algorithm with our theoretical rank bound based on the CIFARIO dataset. For given positive semidefinite matrix, calculate the rank the number of cigenvalues greater than the tolerance 1e 6. The numerical performance is summarized in Table 4. In these cases, one can run our rank reduction algorithm to obtain low-rank solution. Fig. illustrates the procedure by showing the probability mass values associated with the eigenvectors of the estiamted solution S. From the plot, we find that our rank reduction algorithm is capable of producing low-rank solutions even though the matrix size is large.
Iteration 0
Iteration 150
0.0200
of (1 Eigenvalluss
0.0175
Log + Eigenvalues)
x
0.025
0.0125
0.020
0.0100
0.015
0.0075
0.010
0.0050
0.005
0.0025
50
100
150
200
250
300
350
400
100
150
200
250
300
350
400
Support Eigenvalies
Support of Eigenvalues
Iteration 300
Final Iteration
0.07
0.200
Log of Eigenvalues
0.06
0.175
0.05
0.125
0.04
0.100
0.075
0.02
0.050
0.025
200
250
150
350
Support of Eigenvalues
Support of Eigenvalues
Figure 8. Procedure of rank reduction algorithm for CIFARIO example with sample size 200.
L.3. Running Time
27

---

Guarantees of KMS Wasserstein Distances
Table Numerical performance on rank for CIFAR10 dataset
Sample Size
Rank Obtained from
Rank Bound from Theorem 4.5
Algorithm 1
200
400
19
250
500
21
300
600
24
350
700
26
400
800
27
450
900
29
500
1000
31
Table Comparison running time (seconds) different approaches for hypothesis testing, change- and generative modeling experiments.
KMS
Sinkhorn Div
MMD
SW
GSW
MS
ME
CIFAR-10 Testing (n 400)
8.21
4.07
0.017
0.91
2.58
0.064
0.17
CIFAR-10 Testing (n 500)
13.24
4.26
0.021
1.12
2.63
0.073
0.40
CIFAR-10 Testing (n 600)
18.77
4.86
0.026
1.30
2.86
0.083
0.43
CIFAR-10 Testing (n 700)
22.13
5.81
0,033
1.56
2.94
0.096
0.41
CIFAR-10 Testing (n 800)
30.31
6.44
0.040
1.85
3.11
0.110
0.46
CIFAR-10 Testing (n 900)
34.16
7.38
0.047
2.04
3.15
0.120
0.47
CIFAR-I0 Testing (n 1000)
40.18
8.77
0.087
2.28
3.21
0.140
0.52
Human Activity Detection
11.16
26.30
2.42
12.88
22.48
0.58
0.91
Generative Modeling
2530
2212
744
726
1716
790
NA
28


Reviews:
[{"review_id": "rm9e1mesbW", "reviewer": "Reviewer_hSDy", "summary": "This paper establishes statistical theoretical properties of Kernel MSW (K-MSW) distance; the authors provide a finite sample guarantee of K-MSW between empirical probability measures. A performance guarantees are also given when K-MSW is used as a metric of two-sample test. The second part of the paper deals with the computation of 2th power of K-MSW, which is an NP hard problem. The computation relies on a semidefinite relaxation solved with an ineaxact mirroc ascent algorithm.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We very much appreciate reviewer's positive comments and now provide our response on a point-by-point basis:\n\n- [Complexity of KMS?] We would like to clarify that the time complexity of our Algorithm 1 is $\\tilde{O}(n^3\\delta^{-3})$, not $\\tilde{O}(n^2\\delta^{-3})$ as previously stated. In our initial submission, we followed established literature on first-order methods (e.g., [Nemirovsky and Yudin, Problem Complexity and Method Efficiency in Optimization]) which analyzes the complexity of constructing supradient estimators at each iteration. However, we omitted the cost of the proximal gradient projection step in Eq. (13). This step involves computing the matrix exponential and matrix logarithm, each of which requires $O(n^3)$ time. These operations are efficiently implemented in modern linear algebra libraries, but their asymptotic cost remains cubic in $n$.\n\\\n\\\n   Moreover, although Algorithm 1's complexity is independent of the data dimension $d$, it does require precomputing the Gram kernel matrix $G$ (Eq. (8)), and generating the vectors $M_{i,j}'$ as input. This preprocessing step incurs a time complexity of $O(n^3d)$, where the factor of $d$ arises from the kernel computation $k(x,y) = \\exp(-\\|\\|x - y\\|\\|_2^2 / \\sigma^2)$, which scales linearly with $d$ for $x, y \\in \\mathbb{R}^d$.\n\\\n\\\nIn summary, computing the KMS distance requires time complexity $O(n^3(\\delta^{-3} + d))$, which is significantly larger than that of the vanilla sliced Wasserstein distance. We will include a detailed discussion of this point in our revised manuscript.\n\n- [Complexity and Performance Trade-off?] While KMS has in general the highest computational time among the evaluated methods, the overhead is not prohibitive. Importantly, it consistently achieves superior performance, demonstrating that our method is well-suited for practical machine learning applications, albeit with increased computational demands.\n\\\n\\\nPlease also see Table 2 from the Anonymous link https://gofile.io/d/3Pcdxg that reports the numerical running time for hypothesis testing, change-detection, and generative modeling experiments.\nWe observe for the change-detection experiment, our approach even has smaller computational time compared with Sinkhorn Divergence, SW, and GSW.  This efficiency arises because the nonlinear projector can be precomputed using pilot data, enabling fast online computation of the test statistic. In contrast, the baseline methods must recompute the statistics at each detection step, resulting in longer runtimes. We will make this point explicit in the revision to highlight the practical utility of our method.\n\n- [Typos?] We appreciate the reviewer for pointing out our typos. We will correct them in the revision.", "claims_and_evidence": "From a theoretical point view, the authors prove that KMS is a proper metric, and then it has statistical guarantees of finite-sample and two sample as $n$ goes to infinity, the number of support data points. It is worth noting that the convergence rate is dimension-free. KMS is also used as a critical value to determine the rejection of the null hypothesis test between distributions. Concerning the algorithm point, KMS, in a discrete setting, can be cast into a nonconvex max-min optimization problem that is NP-hard (see Theorem 4.2). This problem can be reformulated in a semidefinite relaxation and resolved through an inexact mirror ascent average.", "methods_and_evaluation_criteria": "The proposed method is evaluated on high dimensional hypothesis testing using synthetic and real datasets. Several baselines approach like Sinkhor divergence, MMD, sliced Wasserstein distance (SW), generalized SWD, max-sliced Wasserstein (MS), and optimized mean-embedding test (ME) are tested on CIFRA-10 dataset. The evaluation metric is the testing power where KMS achieves best performance.", "theoretical_claims": "The proofs sound correct.", "experimental_designs_or_analyses": "The numerical experiments are rich and sound good.", "relation_to_broader_scientific_literature": "The proposed approach is included in the optimal transport family metrics used for dimension reduction for high-dimensional data.", "essential_references_not_discussed": "The paper discusses the related works, specifically when comparing the performance with many others approaches for testing hypothesis task.", "other_strengths_and_weaknesses": "- The paper is well written and easy to follow.\n- Proposing statistical guarantees for K-MSW\n- Proving that the 2th power of K-MSW is an NP hard problem.\n- Solving 2-K-MSW through a semidefinite programming relaxation.\n\nWeaknesses:\nKMS suffers from a burden time complexity $\\tilde{O}(n^2d^3)$ compared to the vanilla sliced Wasserstein distance that has $\\tilde{O}(nd)$. This limits its application to several machine learning pipelines, especially generative modeling.", "other_comments_or_suggestions": "The paper is well-written and easy to follow.", "questions_for_authors": "mentioned in weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3oOoPqwRbX", "reviewer": "Reviewer_Av6C", "summary": "1. **Introduction of the Max-Kernel Sliced Wasserstein Distance (KWS)**\n   - The paper presents the Max-Kernel Sliced Wasserstein Distance (KWS), which merges classical max-sliced Optimal Transport (OT) with kernel methods. Data is first mapped to the Reproducing Kernel Hilbert Space (RKHS) through kernel mapping, followed by the computation of max-sliced OT within this space.\n\n2. **Key Properties**\n   - The authors demonstrate several properties of KWS:\n     - Metric property\n     - Sample complexity\n     - Existence of solutions\n\n3. **Computation Details**\n   - The discussion in the paper is primarily on the computation of a special case of KWS, where the cost function is the squared L2 norm (refer to Equation KMS). The paper notes that this KMS is non-convex (see Eq (9)) and NP-hard (refer to Theorem 4.2). The computational approach is outlined in Algorithm 1, including its complexity.\n\n4. **Applications**\n   - KWS is applied in the context of high-dimensional hypothesis testing.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We very much appreciate reviewer's insightful comments and now provide our response on a point-by-point basis:\n\n- [OT in RKHS?] Our formulation is closely related to Zhang et al. (2020), which considers Wasserstein distances between pushforward measures $\\Phi(\\mu)$ and $\\Phi(\\nu)$ via an implicit kernel map $\\Phi$. In Remark 2.6, we show that our KMS distance is equivalent to Max-Sliced Wasserstein in this setting. A key distinction is that our method enjoys sharp statistical convergence rates, which are difficult to obtain in Zhang et al.’s framework due to the curse of dimensionality of Wasserstein.\nMoreover, while their work focuses on kernel embeddings, our motivation differs slightly and is to integrate nonlinear dimensionality reduction with OT. We will clarify this in the revision.\n\n- [OT with a learnable nonlinear map?] We compare our method with $\\mathcal{SW}(\\Phi(\\mu), \\Phi(\\nu))$ where $\\Phi$ is a neural network. Such $\\Phi$ is finite-dimensional and sensitive to neural network architectural choices. In contrast, our KMS can also be reformulated as the similar expression (in Remark 2.6) with $\\Phi$ being an infinite-dimensional, non-parametric kernel mapping. In addition, our method has theoretical guarantee that $\\mathcal{KMS}(\\mu,\\nu)=0$ iff $\\mu=\\nu$. This is important for many applications like two-sample testing, whereas it is not guaranteed in fixed neural network mappings.\n\\\n\\\nWhile we acknowledge that KMS has higher computational cost, it provides a complementary to neural networks. We will incorporate this discussion into the revision and hope the reviewer sees the merits of our proposed framework.\n\n- [Generative Modeling justification?] \nWe revised the setup following [Sliced-Wasserstein Autoencoder] with the new formulation\n$$\n\\min_{\\phi, \\psi} \\mathcal{W}(p_{\\text{data}}, \\psi\\circ\\phi\\circ p_{\\text{data}}) + \\lambda\\cdot\\mathcal{D}(q_{\\text{prior}}, \\phi\\circ p_{\\text{data}}),\n$$\nwhere $\\psi$ and $\\phi$ denote the decoder and encoder, respectively, and $q_{\\text{prior}}$ denotes the pre-defined prior distribution on the latent space (uniform distribution on unit circle). We report experiment results in https://gofile.io/d/3Pcdxg (Table 3).  Our KMS Wasserstein distance has competitive performance as indicated by the smallest Fréchet inception distance (FID) score. Also, it has learnt meaningful latent representation, possibly because it utilizes the flexible kernel-projection mapping to compare data distributions, and thereby achieve competitive performance in generative modeling.\n\\\n\\\nWe also clarify that the experiment setup for this part follows from reference [Sliced-Wasserstein Autoencoder]. It could be possible to improve the performance of all baselines by tuning neural network architectures, optimizers, training time, or even random seed. However, our focus in this paper is theoretical and not to develop state-of-the-art generative models. We hope this experiment will be enough to support the empirical value of KMS.\n\n- [Complexity across all experiments?] The runtime for all methods is reported in the anonymous link (Table 2). While KMS is more expensive, the overhead is not prohibitive.\n\n- [Time Complexity Concerns?] We will correct the stated complexity for solving (11) and add the reference as suggested.\n\n- [Complexity of Algorithm 1?] The reviewer is correct. Time complexity of our Algorithm 1 is $\\tilde{O}(n^3\\delta^{-3})$, not $\\tilde{O}(n^2\\delta^{-3})$ as previously stated. In our initial submission, we followed established literature on first-order methods (e.g., [Nemirovsky and Yudin, Problem Complexity and Method Efficiency in Optimization]) which analyzes the complexity of constructing supgradient estimators at each iteration. However, we omitted the cost of the proximal gradient projection step induced by $h(S)$ in Eq. (13). This step involves computing the matrix exponential and matrix logarithm, each of which requires $O(n^3)$ time. Fortunately, these operations can be efficiently implemented using well-established software packages.\n\n- [Complexity of solving SDR and rank-reduction?] Solving KMS-Wasserstein is NP-hard; we only analyze its convex relaxation. We also refer to our response to Reviewer LEjS for motivation on the rank-reduction algorithm. We acknowledge that further calibrating the optimal solution from SDR to low-rank space is computationally expansive (whose complexity is $\\mathcal{O}(n^5)$), but it is still of theoretical interest if we wish to benefit from the superior performance induced by low-rankness of solution. \n\n- [Typos?] We thank the reviewer and will correct all noted typos and clarify notations in the revision.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes.", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "Classical Optimal Transport (OT) literature often examines sliced optimal transport and generalized sliced optimal transport within the dataset space. This paper introduces an innovative combination by integrating the kernel method with sliced OT, effectively extending sliced OT into the Reproducing Kernel Hilbert Space (RKHS). Thus, the methodology presented fundamentally operates as sliced OT within RKHS.", "essential_references_not_discussed": "Yes. \n\nThe following paper also discusses the OT in RKHS space. \n\nZhang, Z., Wang, M., & Nehorai, A. (2020). Optimal Transport in Reproducing Kernel Hilbert Spaces: Theory and Applications. *IEEE Transactions on Pattern Analysis and Machine Intelligence*, 42(7). [DOI: 10.1109/TPAMI.2019.2903050](https://doi.org/10.1109/TPAMI.2019.2903050)", "other_strengths_and_weaknesses": "**Weaknesses:**\n\n1. **Advantages of Kernel Method and OT:**\n   - The authors are encouraged to further emphasize the benefits of integrating the kernel method with Optimal Transport (OT). Given the paper's proposal of kernel sliced OT (KMS) and its inherent complexities, it is crucial to justify the advantages of this approach over simpler alternatives like applying a non-linear learnable mapping (e.g., MLP) followed by sliced OT. Particularly, the substantial disadvantages of KMS—its NP-hard nature and the expensive, approximative solver—demand a strong argument for using the kernel method instead of just non-linear mapping with classical OT.\n\n2. **Generalized Sliced Wasserstein in Modeling:**\n   - Including Generalized Sliced Wasserstein (GSW) in generative modeling experiments would provide a comprehensive evaluation.\n   - A comparison of wall-clock times and data sizes across all experiments should be added to assess the practical implementation of the discussed methods.\n\n3. **Time Complexity Concerns:**\n   - The paper's stated complexity for solving the inner optimization problem (11) as \\(\\tilde{\\mathcal{O}}(n^2/\\epsilon)\\) appears to be understated. Standard notation would suggest \\(\\mathcal{O}(n^2 \\ln(n)/\\epsilon)\\), as supported by recent research [1].\n   - The complexity of eigen decomposition required for computing \\(h(S)\\) is \\(O(n^3)\\), yet this does not seem to be included in the time complexity discussed in Theorem 4.4.\n   - The precision gap (\\(\\delta\\)) addressed in Theorem 4.4 seems to refer to the SDR problem's accuracy rather than the original KMS problem. According to Theorem 4.6, the computational complexity for an approximate solution of KMS is projected at \\(n^5\\), raising further concerns about its practicality.\n\n4. **Unexplained Superior Performance in Generative Modeling:**\n   - There is a lack of explanation for why KMS yields better performance in the generative model experiments. Providing a detailed theoretical justification or analysis would greatly enhance the credibility and acceptance of KMS's claimed effectiveness.\n\n**Reference:**\n[1] D. Dvinskikh and D. Tiapkin, \"Improved complexity bounds in Wasserstein barycenter problem,\" in *International Conference on Artificial Intelligence and Statistics*, PMLR, 2021, pp. 1738–1746.", "other_comments_or_suggestions": "1. **Clarification Needed on Notation**\n   - In Equation (9), the notation \\(M\\) is introduced without explanation. It appears to reference \\(M = G - G^T\\) from Equation (7). Clarification in the manuscript would be beneficial.\n\n2. **Typographical Corrections**\n   - Page 4, line 202, correct the notation from \"O(n^1(/2p))\" to \"O(n^{1/2p})\".\n   - Page 5, update \"\\(\\pi^*(S) \\in \\Gamma(S)\\)\" to \"\\(\\pi^*(S) \\in \\Gamma_n(S)\\)\" to ensure accurate mathematical representation.", "questions_for_authors": "See the weakness part.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "eP9h3izuuP", "reviewer": "Reviewer_LEjS", "summary": "This paper studies statistical and computational properties of Kernel Max-Sliced Wasserstein distances. On the statistical side, the paper's main result is a high-probability bound on the KMS Wasserstein distance between a distribution and the empirical distribution of samples from that distribution. This is in turn used to justify the cutoff for a nonparametric two-sample test, whose power is lower bounded assuming the KMS Wasserstein distance between the true distributions is at least this threshold. On the computational side, although prior work used a representer theorem to reduce computation of the KMS Wasserstein distance between two empirical distributions to a finite-dimensional optimization problem, the present paper shows that this finite-dimensional problem is NP-hard in the worst case. Therefore, the paper proposes a semidefinite relaxation of this problem, which has complexity polynomial in the sample size, precision of the solution, and a certain norm of the Gram matrices of the data. The paper also bounds the rank of the true solution to the semidefinite relaxation and suggests a rank-reduction algorithm to produce solutions with rank near that of the true solution. Finally, the paper provides experimental results demonstrating the computational advantages of the proposed algorithms as well as strong performance on a variety of two-sample testing problems, and applications to change-point detection and generative modeling.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer's positive comments and provide our response below:\n \n- [relationship and distinctions with literature [9]?] We appreciate the reviewer for highlighting this point. Literature [9] establishes the statistical convergence rate for the empirical Max-Sliced (MS) distance as $O(R \\cdot n^{-1/(2p)})$, where $R$ denotes the diameter of the sample space and $n$ is the sample size. This rate is minimax optimal, and the compactness assumption on the sample space is crucial to the analysis. In contrast, our results show that under the bounded kernel assumption (Assumption 1), the empirical KMS distance achieves the same convergence rate of $O(n^{-1/(2p)})$, without requiring the compactness of the original sample space. This allows our result to apply to a broader class of probability distributions. Furthermore, we prove that this rate is also minimax optimal.\n\\\n\\\nOur statistical analysis builds on the insights from [9]. Specifically, the KMS distance can be interpreted as first mapping the data distributions into an infinite-dimensional Hilbert space via an implicit feature map (induced by the kernel), and then applying the MS distance in that space (see our discussion in Remark 2.6). A key step in our proof leverages the statistical results for the MS distance in Hilbert spaces from [9, Corollary 2.8]. Our bounded kernel assumption ensures that the transformed distributions have bounded support (i.e., finite diameter) in the Hilbert space. We will include this discussion in our revision.\n- [discussion of decreasing power issue of KMS?] We appreciate the reviewer for this insightful comment. Under fair alternative, we believe that the power of our KMS distance will decrease, as the strength of our assumption $\\mathrm{KMS}(\\mu,\\nu) - \\Delta(n,\\alpha)>0$ increases. We will emphasize this point in our revision and add an additional experiment to illustrate the trend of decreasing power of our method compared to other baselines under such settings.\n- [minimax lower bound?] We provide the following example to demonstrate that the bound $\\Delta(n, \\alpha)$ is tight. Consider the case where $\\mathcal{B} = [0,1]$, the kernel is $k(x,y) = xy$, and the distribution is $\\mu = \\frac{1}{2}\\delta_0 + \\frac{1}{2}\\delta_1$. Then the empirical KMS distance is given by $\\mathrm{KMS}(\\mu, \\hat{\\mu}_n) = \\left|\\frac{1}{2} - \\frac{N}{n}\\right|^{-1/p}$, where $N \\sim \\mathrm{Binom}(n, \\frac{1}{2})$. It can be shown that $\\mathbb{E}[\\mathrm{KMS}(\\mu, \\hat{\\mu}_n)] = \\Theta(n^{-1/(2p)})$, thereby confirming that the rate $O(n^{-1/(2p)})$ is optimal in the worst case.\n- [Motivation for rank analysis?] Our motivation for studying the rank of the solution is that it has impact on the quality of the resulting approximation to the original non-convex optimization problem.\nRecall that globally solving the KMS Wassertein distance involves a rank-1 constraint, which our SDR formulation relaxes it.\nWhen the solution to SDR is of rank higher than $1$, we resort to constructing a feasible rank-1 solution by taking its leading eigenvector. However, if the SDR solution is already low-rank, then the rounding procedure yields a solution that is closer to the ground-truth rank-1 optimum. \n\\\n\\\nWe include an experimental study (see Figure 6 from the Anonymous link https://gofile.io/d/3Pcdxg) that visualizes how our rank reduction algorithm gradually reduces the rank of the SDR solution.\nOriginally the optimal solution to SDR is of rank $400$, and our rank reduction algorithm iteratively reduces its rank by $1$ until we obtain the $19$-rank solution. We show the quality of the rounded feasible solution to the original KMS Wasserstein distance problem, and observe the corresponding objective will increase from $0.63$ to $0.68$.\n\n\n- [Typos?] We appreciate the reviewer for pointing out our typos. We will correct them in the revision.", "claims_and_evidence": "The claims are well justified.", "methods_and_evaluation_criteria": "Although not the main contribution of the paper, the empirical evaluations effectively demonstrate the utility of the theoretical results and proposed algorithms.", "theoretical_claims": "I skimmed the proof of Theorem 3.2, which seemed reasonable.", "experimental_designs_or_analyses": "The experiments seem extensive, although I did not check the details.", "relation_to_broader_scientific_literature": "1) The relationship/distinctions between the statistical results in the current paper and those in [9] are not very easy to understand from the brief description under Related Work. Can the authors elaborate on what the current statistical results add beyond those of [9]?\n2) Corollary 3.3 and Remark 3.4: I think it's worth adding a comment distinguishing dimension-independence of the problem of *estimating KMS Wasserstein distance* from the problem of *two-sample testing using KMS Wasserstein distance*. Under \"fair\" alternatives (see Ramdas et al. 2015; full reference below), the true KMS Wasserstein distance probably decreases, i.e., the strength of the assumption KMSp(µ, ν) − ∆(n, α) > 0 increases, with dimension. See Ramdas et al. (2015; reference below) for detailed discussion of this phenomenon in the case of MMD.\n\n**References**\nRamdas, A., Reddi, S. J., Póczos, B., Singh, A., & Wasserman, L. (2015, March). On the decreasing power of kernel and distance based nonparametric hypothesis tests in high dimensions. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 29, No. 1).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "- Example 2.5 is quite nice for illustrating the advantages of KMS-Wasserstein distance over MS-Wasserstein distance.\n- Generally, the paper is clearly written and easy to read.", "other_comments_or_suggestions": "- Typo: P. 14, lines 761-769: This inequality is missing a $\\mathcal{KMS}_p(\\mu, \\nu)$ term.\n\n- Theorem 2.4 should probably also state that $\\mathcal{KMS}_p$ satisfies the triangle inequality, since this is used in the proof of Theorem 3.2, Part II.\n\n- Typo: Line 193, Column 1: \"Wassersrein\" -> \"Wasserstein\"", "questions_for_authors": "1) Is it possible to provide a minimax lower bound, or at least an example, showing that the rate $\\Delta(n, \\alpha)$ in Theorem 3.2 is tight?\n2) I don't quite get the motivation for analyzing the rank of the SDR solution. I understand that solutions to the original problem should be rank-1, but what are the ramifications of producing, say, a rank-2 solution vs a rank-3 solution?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The focus of the submission is the analysis of the kernel max-sliced (KMS) $p$-Wasserstein distance (Def. 2.3; $1\le p<\infty$). Particularly, the authors
1) prove (sharp) finite-sample guarantee for the KMS $p$-Wasserstein distance between the distribution and its empirical variant (Theorem 3.2),
2) show that the exact computation of the KMS $p$-Wasserstein for $p=2$ is NP-hard (Theorem 4.2), and propose a semidefinite relaxation with complexity (Theorem 4.4) and rank bounds (Theorem 4.5). The efficiency of the proposed approach is demonstrated on both synthetic and real-world experiments including the detection of human activity transition and generative modelling.

Estimating information theoretical measures is a key problem in modern data science with various successful applications. The authors provide tools in this context which have important theoretical and numerical implications, in a well-organized and clearly-written work as it was assessed by the reviewers.
Decision: Accept (poster)