Wasserstein $K$-means for clustering probability distributions
Authors: Yubo Zhuang
Xiaohui Chen
Yun Yang
Venue: NeurIPS 2022 Accept (2022)
Keywords: Barycenter
K-means clustering
Optimal transport
Semi-definite programming
Wasserstein distance

Abstract: Clustering is an important exploratory data analysis technique to group objects based on their similarity. The widely used $K$-means clustering method relies on some notion of distance to partition data into a fewer number of groups. In the Euclidean space, centroid-based and distance-based formulations of the $K$-means are equivalent. In modern machine learning applications, data often arise as probability distributions and a natural generalization to handle measure-valued data is to use the optimal transport metric. Due to non-negative Alexandrov curvature of the Wasserstein space, barycenters suffer from regularity and non-robustness issues. The peculiar behaviors of Wasserstein barycenters may make the centroid-based formulation fail to represent the within-cluster data points, while the more direct distance-based $K$-means approach and its semidefinite program (SDP) relaxation are capable of recovering the true cluster labels. In the special case of clustering Gaussian distributions, we show that the SDP relaxed Wasserstein $K$-means can achieve exact recovery given the clusters are well-separated under the $2$-Wasserstein metric. Our simulation and real data examples also demonstrate that distance-based $K$-means can achieve better classification performance over the standard centroid-based $K$-means for clustering probability distributions and images.

Wasserstein K-means for clustering probability distributions
Yubo Zhuang, Xiaohui Chen, Yun Yang Department of Statistics
University of Illinois at Urbana- Champaagn {yubo2, xhchen, 0allle@illliooo edu
Abstract
Clustering is important exploratory data analysis technique to group objects based on their similarity. The widely used K-means clustering method relies some notion of distance partition data into a fewer number of groups. In the Euclidean space, centroid-based and distance-based formulations of the K- means equivalent. In modern machine learning applications, data often arise as probability distributions and natural generalization to handle measure-valued data is the optimal transport metric. Due to non-negative Alexandrov curvature of the Wasserstein space, barycenters suffer from regularity and non- robustness issues. The peculiar behaviors of Wasserstein barycenters may make the centroid-basee formulation fail to represent the within-cluster data points, while the more direct distance-based K-means approach and its semidefinite program (SDP) relaxation are capable of recovering the true cluster labels. In the special case clustering Gaussian distributions, we show that the SDP relaxed Wasserstein K K-means can achieve exact recovery given the clusters are well-separated under the 2- Wasserstein metric. Our simulation and real data examples also demonstrate that distance- K-means can achieve better classification performance over the standard centroid-based K-means for clustering probability distributions and images.
Introduction
Clustering is a major tool for unsupervised machine learning problems and exploratory data analysis in statistics. Suppose we observe a sample of data points X XX, taking values in a metric space (X', || II). Suppose there exists clustering structure G Gi...., G such that each data point X, belongs to exactly one of the unknown cluster G The goal of clustering analysis to recover the true clusters G G given the input data XI...., X,, In the Euclidean space X RP, the K-means clustering is a widely used method that achieves the empirical success in many applications [MacQueen, 1967]. In modern machine learning and data science problems such as computer graphics [Solomon et al., 2015], data exhibits complex geometric features and traditional clustering methods developed for Euclidean data may not be well suited to analyze such data.
In this paper, we consider the clustering problem of probability measures p1..... Pn into K groups. As motivating example, the MNIST dataset contains images of handwritten digits 0-9. Normalizing the greyscale images into histograms as probability measures, a common task is to cluster the images. One can certainly apply the Fuclidean K-means to the vectorized images. However, this would lose important geometric information the two-dimensional data. the other hand, theory of optimal transport [Villani, 2003] provides appealing framework to model measure-valued data as probabilities in many statistical tasks [Domazakis et al., 2019, Chen et al., 2021, Bigot et al., 2017, Seguy and Cuturi, 2015, Rigollet and Weed, 2019, Hütter and Rigollet, 2019, Cazelles al., 2018).
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Background on K-means clustering. Algorithmically. the K --means clustering have two equivalent formulations in the Euclidean space centroid based and iistance-based the sense that they both yield the same partition estimate for the true clustering structure. Given the Euclidean data X1,..., X E IR", the centroid-based formulation of standard K-means can be expressed as
K
K
min P1... PK ER
min IX - B&112 kE[K]
min G Gi
| XAli
G
Gx
In]
(I)
Rm IEG
where clusters {G&] are determined by the Voronoi diagram from {B& Xx
GxI-->EeGs X; denotes the centroid of cluster Gg. L denotes the disjoint union and [n] [L....,n]. Heuristic algorithm for solving (1) includes Lloyd's algorithm [Lloyd, 1982], which is an iterative procedure alternating the partition and centroid estimation steps. Specifically, given an first assigns each data point its nearest centroid the I-th iteration according to the Voronoi diagram, i.e.,
initial centroid estimate 8 ..... ,,KK _kk
G = i [m] |X - 11 & 11, - 1/212. vj E [1]]]
-8,
(2)
and then update the centroid for each cluster
5 s(++1)
Σ X
(3)
|G iEG
where G denotes the cardinality of G Step (2) and step (3) alternate until convergence.
The distance-batt based (sometimes also referred partition-based) formulation directly solves the following constrained optimization problem without referring to the estimated centroids:
K
K G+ = |m]]
min GI.....C
Σ |X, L.EG
-X,-x111
(4)
GR
A=
Observe that (1) with nearest centroid assignment and (4) are equivalent for the clustering purpose due to the following identity, which extends the parallelogram law from two points to n points,
n Σ -X,-,1
2n Σ ||
X12. with x
n -xx:
X,
and X, R..
(5)
Consequently, the two criteria yield the same partition estimate for Gi.... ... The key identity (5) establishing the equivalence relies two facts of the Euclidean space: it vector space (i.c.. vectors can be averaged in the linear sense); (ii) it is flat (i.c., zero curvature), both of which are unfortunately not true for the Wasserstein space (P2 (R?), W2) that endows the space P2 (2(I)) of all probability distributions with finite second moments with the 2- Wasserstein metric W2 [Ambrosio al., 2005]. particular, the 2- Wassersteii distance between two distributions. μ and (2(P defined as
W (14, v) = min
> vll2 (2x
(6)
xB*
where minimization over runs over all possible couplings with marginals l and v. It is well-known that the Wasserstein space metric space (in fact geodesic space) with non-negative curvature in the Alexandrov sense [Lott, 2008].
Our contributions. We summarize main contributions as followings: (i) we provide evidence for pitfalls (irregularity and non-robustness) of barycenter- Wasserstein K-means, both theoretically and empirically, and (ii) generalize the distance-based formulation of K-means to the Wasserstein space and establish the exact recovery property of its SDP relaxation for clustering Gaussian measures under separateness lower bound in the 2- Wasserstein distance.
Existing work. Since the K-means clustering is a worst-case NP-hard problem [Aloise et al., 2009], approximation algorithms have been extensively studied in literature including: Llyod's algorithm [Iloyd, 1982), spectral methods [von Luxburg, 2007, Meila and Shi, 2001, Ng al., 2001], semidefinite programming (SDP) relaxations [Peng and Wei, 2007], non-convex methods via

---

low-rank matrix factorization [Burer and Monteiro, 2003]. Theoretic guarantees of those methods are established for statistical models on Euclidean data [Lu and Zhou, 2016, von Luxburg et al., 2008, Vempala and Wang. 2004, Fei and Chen, 2018, Giraud and Verzelen, 2018, Chen and Yang, 2021, Zhuang et al., 2022].
The concept of clustering general measure-valued data is introduced by Domazakis et al. [2019], where the authors proposed the centroid-based Wasserstein K-means algorithm. It replaced the Euclidean norm and sample means by the Wasserstein distance and barycenters respectively. Verdinelli and Wasserman [2019] proposed a modified Wasserstein distance for distribution clustering. And after that, Chazal et al. [2021] proposed a method in Clustering of measures via mean measure quantization by first vectorizing the measures in a finite Euclidean space followed by an efficient clustering algorithm such as single-linkage clustering with L distance. The vectorization methods could improve the computational efficiency but might not capture the properties of probability measures well compared to clustering algorithms based directly on Wasserstein space.
Wasserstein K-means clustering methods
In this section, we generalize the Euclidean K-means to the Wasserstein space. Our starting point is to mimic the standard K-means methods for Euclidean data. Thus we may define two versions of the Wasserstein K- K-means clustering formulations: centroid-busedd and distance-bused. As we mentioned in Section 1, when working with Wasserstein space (P2(P)) W2), the corresponding centroid-based criterion (1) and the distance-based criterion (4), where the Euclidean metric || |/2 is replaced with the 2- Wasserstein metric W2, may lead to radically different clustering schemes. To begin with, we would like to argue that due to the irregularity and con-robbst of barycenters in the Wasserstein space, the centroid- based criterion may lead to unreasonable clustering schemes that lack physical interpretations and are sensitive to small data perturbations.
2.1 Clustering based barycenters
The centroid-based Wasserstein K-means for extending the Lloyd's algorithm into the Wasserstein space has been recently considered by Domazakis et al. [2019]. Specifically, it is an iterative 1)
algorithm proceeds as following. Given an initial centroid estimate (1)) ..... "x one first assigns
each probability measure (4,....; Mrn its nearest centroid in the Wasserstein geometry at the f-th iteration according to the Voronoi diagram:
c
(t) = i E [m] Wa(m,, (0) & Wa(mn, Vj E
and then update the centroid for each cluster
(t+1) arg min (E (EG
Σ (2
(8)
Note that v/ ((+1) in (8) is referred as barycenter of probability measures J4, G a Wasserstein analog of the Euclidean average or mean [Agueh and Carlier, 2011]. We will also ex-changeably use barycenter-based K-means to mean the centroid-based K-means in the Wasserstein space. Even though the Wasserstein barycenter is natural notion of averaging probability measures, it may exhibit peculiar behaviours and fail to represent the within-cluster data points, partly due to the violation of the generalized parallelogram law (5) (for non-flat spaces) the Euclidean metric || ll2 is replaced with the 2- Wasserstein metric W2.
Example (Irregularity of Wasserstein barycenters). Wasserstein barycenter has much less regu- larity than the sample mean in the Euclidean space [Kim and Pass, 2017]. In particular, Santambrogio and Wang [2016] constructed a simple example of two probability measures that are supported on line segments in R²? whereas the support of their barycenter obtained as the displacement interpolation the two endpoint probability measures is not convex (cf. left plot in Figure 1). In this example, the probability density 840 and μy are supported on the line segments Lo = {(x,as)= as)
[-1, 1]) and T, = ((x, -as) 8 [-1, -1,1] 1]} respectively. We choose a E (0, 1) to identify the orientation of T.o and L1 based on the a-axis. Moreover, we consider the lincar density functions (o(s) (1 =2//2
3

---

and (s) = (1 + s) / for [-1,1] supported on Lo and L1 respectively. Then the optimal transport map T := T po- from AO to μl given by
T(z, (-1+ - -a- (-1+ (1-11)33) z)3))
and the barycenter corresponds to the displacement interpolation = [(1 1)id + 1 at t 0.5 [McCann, 1997]. For self-contained purpose, we give the proof of (9) in Appendix D.I. Fig. on the left shows the support of barycenter boundary) even though the supports of μo and μl are convex. This example shows that the barycenter functional is not geodesically convex in the Wasserstein space. As barycenters turn out to be essential in centroid-base Wasserstein K K-mmans and irregularity of the barycenter may fail represent the cluster (see more details in Example and Remark 9 below), this counter-exampae is our motivation seek alternative formulation.
is not convex (in fact part of an ellipse
P. -0.5
-1.0
-0.5
0.0
0.5
1.0 -1.0
-0.5
0.0
0.5
3.0
1.5
2.0
Figure 1: Left: support of the Wasserstein barycenter as the displacement interpolation between μo and μl at 0.5 in Example Right: con-robustness of the optimal transport map (arrow lines) and Wasserstein barycenter w.r.t. small perturbation around a = 1 for the target measure in Example 2.
Example 2 (Non-robustness of Wasserstein barycenters). Another unappealing feature of the Wasserstein barycenter is its sensitivity to data perturbation: a small (local) change in one contributing probability measure may lead to large (global) changes in the resulting barycenter. See Fig. on the right for such an example. 0.5 -..-1) 1) and the target measure μl 0.5 (-5, -,)) + 0.5 (,5(11)) for some 0. It is easy that the optimal transport map T = T has dichotomy behavior:
this example, we take the source measure as μAD = 0.5 0.5 ((1.1
{
ir (-) (1, ifa>
and T(1, = { -
(L.a) ir
T(-1,1)
(-).-)) ifa>
(10)
Thus the Wasserstein barycenter determined by the displacement interpolation μt [(1 ((1-t)it ((n-t)id+TTlp44 is discontinuous function at a 1. This non-robustness can be attribute the discontinuity of the Wasserstein barycenter as function of its input probability measures; in contrast, the Euclidean mean is a globally Lipchitz continuous function of its input points.
Because of these pitfalls of the Wasserstein barycenter shown in Examples and 2, the centroid- based Wasserstein K-means approach described at the beginning of this subsection may lead to unreasonable and unstable clustering schemes. In addition, ill-conditioned configuration may significantly slow down the convergence of commonly used barycenter approximating algorithms such iterative Bregman projections [Benamou et al., 2015]. Below, give concrete example of such phenomenon in the clustering context.
Example 3 (Failure of centroid-bri based Wasserstein K K-means). In nutshell, the failure in this example due to the counter-intuitive phenomenon illustrated in the right panel of Fig. 2, where some distribution μ4 in the Wasserstein space may have larger W2 distance Wasserstein barycenter μi than every distribution (4; = 1, 2) that together forms it. As result of this strange configuration, even though μy is closer μr and μl from the first cluster with barycenter μ than μ4 coming from a second cluster with barycenter μ4z it will be incorrectly assigned to the second cluster using the centroid-based criterion (7), since W2(u3, 24)= > W2(53, (5)) > max { (W2 f1), W2(13, (2))- In contrast, for Euclidean spaces due to the following equivalent formulation of the generalized

---

parallelogram law (5),
Σ | - = m|| x
+
|X 111 3 n||x x13. for any X E R2,
==1
there is always some point X satisfying - || - 1112. that is, further away from X
than the mean X; thereby excluding counter-intuitive phenomena the one shown in Fig. 2.
12
111
P4
12
J41
μ4 = u2
P2
J13
13
Figure 2: Left: visualization of Example in 2² and Wasserstein space. Right: the black curve connecting ja1 and μl2 depicts the geodesic between them.
Concretely, the first cluster Gj is shown the left panel of Fig. 2 highlighted by red circle, consisting of m copies of (H1, (22) pairs and one copy of p3: the second cluster G containing copies of μ4 is highlighted by blue circle. Each distribution assigns equal probability mass two points, where the two supporting points are connected by a dashed line for easy illustration. More specifically, we set
A1 = 0.3 ((..)) + 0.5 00,56(-z.-0))) μ22 = 0.3 (sr- + 0.5 6.556 ((2.)))
[43 = 0.5 ((++e..0) + 0.5 (.5t and
14 = 0.5 (.55((tttttea..))) + 0.5 6( ((e-et+a.-v)
where ((m.v) denotes the point mass measure at point (x, v), and (z, V, é1, 62) are positive constants.
The property of this configuration can summarized by the following lemma.
Lemma (Configuration characterization) If (Ir, (I, 61, satisfies
y² < min 2², (025 A., and An. < o < A. +
where A == e 2x + 2xc1, then for all sufficiently large (number of copies of μr and p2).
W2(H3, < W2 Mi)
and
max max W2(ui, k-1,21,EGG
min W2(ue,
GEG,JGGZ
largest within-cluster distance
least between-cluster distance
where 1t denotes the Wasserstein barycenter of cluster G3 for = 1,2.
Note that the condition of Lemma 4 implies y æ. Therefore, the barycenter between μ4 andu is 7 == 0.5 ( + (-z.) lying on the horizontal axis. By increasing m, the barycenter μ of cluster Gi can be made arbitrarily close j²* The second inequality in Lemma shows that all within-cluster distances are strictly less than the between-cluster distances; therefore, clustering based pairwise distances is able correctly recover the cluster label of μ4z. However, since μ3 is closer to the barycenter 122 of cluster G according to the first inequality in Lemma 4, it will be mis-classified into G using the centroid-based criterion. We emphasize that cluster positions in this example are generic and do exist in real data; see Remark 9 and Section 4.3 for further discussions on our experiment results on MNIST data. Moreover, similar Example change in the orientation of distribution μl may completely alter the clustering membership of μ4 based on the centroid criterion. Specifically, if we slightly increase to make it exceed y, then the barycenter between μl and μl becomes Ø == 0.5 850.;) + 0.5 8 that lies on the vertical axis. Correspondingly, if based on centroids, then μ3 should be clustered into Gi as is closer to the barycenter μe of Gi than the barycenter μī of G Therefore, the centroid- based criterion can be unstable against data perturbations. In comparison, a pairwise distances based criterion always assigns #3 into cluster G matter y or y.
small
5

---

2.2
Clustering based on pairwise distances
Due to the irregularity and non non-robustness of centroid-based Wasserstein K -eams,, we instead propose and advocate the use of distance-based Wasserstein K-means below, which extends the Euclidean distance-basee based K X-mmans formulation (4) into the Wasserstein space,
K G [m]
min Gi....gx
Σ i.ECC,
My): (
}
(11)
GR
k=1
k=1
Correspondingly, we can analogously design a greedy algorithm resembling the Wasserstein Lloyd's algorithm described in Section 2.1 that solves the centroid-basee Wasserstein K-means. Specifically, the greedy algorithm proceeds in an iterative manner as following. Given an initial cluster membership estimate G Gkk one assigns each probability measure /1,....: Mn based minimizing the averaged squared W2 distances to all current members in every cluster, leading to an updated cluster membership estimate
Gl = fi i
-iiE
[n]:
[n]:
(2(, Σ S
W2(, 1G SEG
Σ
Vj [kI]. (12)
G
We arbitrarily select among the least W2 distance clusters in the case of a tie. We highlight that the center-based and distance-based Wasserstein K-means formulations may not necessarily be equivalent to yield the same cluster labels (cf.. Example 3). Below, we shall give some example illustrating connections to the standard K --mans clustering in the Euclidean space.
Example 5 (Degenerate probability measures). If the probability measures are Dirac at point X, E RP, i.e., P4 5x,, then the Wasserstein K-means is the same as the standard K-means since W2(pi, = X Xjll2-
Example (Gaussian measures). If μ4, = N (m, Vi) with positive-defintt covariance matrices H, > O, then the squared 2- 2-Wasserstein distance can be expressed as the sum of the squared Euclidean distance on the mean vector and
1/2
(V.VV, Tr V -
(13)
the squared Bures distance the covariance matrix [Bhatia al., 2019]. Here, we use 1//2 to denote the unique symmetric square root matrix of V > 0. That is,
(24. = [m; - m31 + d (V.V, V2).
(14)
Then the Wasserstein K-means, formulated either in (7) (11), can be viewed covariance- adjusted Euclidean K-means by taking account into the shape or orientation information in the -non-degenerate) Gaussian inputs.
Example 7 OOne-dimensional probability measures). If μ4 are probability measures on K with cumulative distribution function (cdf) F;, then the Wasserstein distance can be written in terms the quantile transform
W (1s, (1) ( (F (1)
(1)) du,
(15)
where F- the generalized inverse of the edf F on J0, 1] defined (w) e R F(x) > u (cf. Theorem 2.18 [Villani, 2003J). Thus the one-dimensional probability measures in Wasserstein space can be isometrically embedded in flat L2 space, and we bring back the equivalence of the Wasserstein and Euclidean K-means clustering methods.
3 SDP relaxation and its theoretic guarantee
Note that Wasserstein Lloyd's algorithm requires to use and compute the barycenter (7) and (8) at each iteration, which can be computationally expensive when the domain dimension d is large or the configuration is ill-conditioned (cf. Example 2). On the other hand, known that solving the distance-based K K-means (4) is worst-case NP-hard for Euclidean data. Thus we expect solving

---

the distance- Distance-based Wasserstein K- K-means (11) is also computationally hard. A common way is to consider convex relaxations to approximate the solution of (11). It known that certain SDP relaxation is iiformation-theortticcll
tight for (4) when the data XI...,, X, E RR are generated from a Gaussian mixture model with isotropic known variance [Chen and Yang. 2021]. this paper, we extend the idea into Wasserstein setting for solving (11).
typical SDP relaxation for Euclidean data uses pairwise inner products to construct an affinity matrix for clustering [Peng and Wei, 2007]; unfortunately, due to the non-flatness nature, a globally well- defined inner product does not exist for Wasserstein spaces with dimension higher than one. Therefore, we will derive a Wasserstein SDP relaxation to the combinatorial optimization problem (4) using the squared distance matrix Anx= = {aty] with aij W (14, Mt). Concretely, we can one-to-one reparameterize any partition (G1,.. Gg) as a binary assignment matrix H { {0, 1) 1
such that fasa = 1 if E G5 and hy = 0 otherwise. Then (11) can be expressed as nonlinear 0-1 integer program,
min (AAH H E 10, H1K = 1n).
(16)
where I is the X 1 vector of all ones and B diag( |1| 16,1-,,,..,, ...., .....,GGK|-y
Changing of variable
to the membership matrix Z = HBHTT we note that Znxn xn is symmetric positive semidefinite (psd) matrix Z -/2 O such that Tr(Z)= = K, 21, = I,, and Z > 0 entrywise. Thus we obtain the SDP relaxation of (11) by only preserving these convex constraints:
(A, TT = Z, Z 0. Tr(2) = K, 21, = 1,: 2 >>>00 oJ.
min ZER
(17)
To theoretically justify the SDP formulation (17) of Wasserstein K-means, we consider the scenario
of clustering Gaussian distributions in Example 6, where the Wasserstein distance (14) contains two separate components: the Euclidean distance mean vector and the Bures distance (13) on covariance matrix. Without loss of generality, we focus on mean-zero Gaussian distributions since optimal separation conditions for exact recovery based on the Euclidean mean component have been established in [Chen and Yang. 2021]. Suppose we observe Gaussian distributions V N(0, Vi), i [n] from K groups Gi. Gkk where cluster G7 contains ni members, and the covariance matrices have the following clustering structure: i E Gj. then
V = (I + +xx, (I+ tx,, with XI...., Xn f.ddd SymNN0.,
(18)
where the psd matrix (+) is the center of the k-th cluster, SymN 1) denotes the symmetric random matrix with i.i.d. standard normal entries, and t is small perturbation parameter such that ++x)),) is
psd with high probability. For zero-mean Gaussian distributions, we have W2(N (0, V), (N,0, U)) U(V,U) according to (14). Note that on the Riemannian manifold of psd matrices, the geodesic emanating from (R) in the direction X as a symmetric matrix can be linearized by V = (T tx ( in small neighborhood of t, thus motivating the parameterization of our statistical model in (18). The next theorem gives a separation lower bound to ensure exact recovery of the clustering labels for Gaussian distributions.
Theorem (Exact recovery for clustering Gaussians). Let A == mink (V VO) denote the minimal pairwise separation among clusters, := maxkg [K] 1l (and n =- minke[[] ni)) the maximum (minimum) cluster size, and m : min 2ng,n the minimal pairwise harmonic mean of cluster sizes. Suppose the covariance matrix V; of Gaussian distribution v; = No, V,) is independently drawn from model (18) for i 1, 2,...,, n. Let B (0, If the separation A satisfies
tut
C₁2
A A min (00((11,,,993 32
V Dp" logn,
(19)
then the SDP (17) achieves exact recovery with probability at least 1 C2n-1, provided that
n C3 log 2 n. 1 5 Caylogni n/ [(p + log n)) / n/m Cs log n,
where == max;
||(]]] up T. max Tr [r[(vV (v()) J and C., i 1,2, 3, 4,5 are constants.
Remark 9 (Further insight on pitfalls of barycenter-basee Wasserstein K-means). Theorem 8 suggests that different from Euclidean data, distributions after centering can be clustered if scales and rotation angles vary (i.e., covariancc-adjjsted)). We further illustrate the rotation and scale effects on
7

---

a²
Ho
112
Mo
5
14
S
i
Mo
112
p u Baryoenters from B-WKM.
dc 0,0561
0
d = 0.0383
S
Pi-P Trie baryoemnens for Gj. G : One image from cluster G7-
Figure 3: Visualization of misclassification for the barycenter- based Wasserstein K-means (B- WKM) on a randomly sampled subset from MNIST (200 digit "0" and 100 digit "5"). The plot at the bottom a example misclassified image. The right plot the abstraction of the images the Wasserstein space. The color depth indicates the frequency of the distributions. Red and blue colors stand for distributions belong to true clusters "0" and "5".
the MNIST data that may mislead the centroid bested Wasserstein K K-means, thus providing a real data support for Example Here we randomly sample two unbalanced clusters with 200 numbers of "O" and 100 numbers of "5". Fig. 3 shows the clustering results for the centroid-based Wasserstein K-means and its oracle version where we replace the estimated barycenters (11 A22 with the true barycenters ML, μC computed on the true labels. Comparing the Wasserstein distances W2(Ho, mi) and W2(Ho. 145), we see that the image μo (containing digit "0") is closer to μe (true barycenter of digit "5") and thus it cannot be classified correctly based on the nearest true barycenter (cf. Fig. on the left). Moreover, Wasserstein K-means based on estimated barycenters μt, Hlz yields two clusters of mixed "0" and "5". In both cases, the misclassification error is characterized by grouping similar degrees of angle and/or stretch. Since there are two highly unbalanced clusters of distributions, Wasserstein K-means is likely to enforce larger cluster to separate into clusters and absorb those around centers (cf. Fig. 3 on the right), leading to larger classification errors. We shall see that in Section 4.3 the distance-baid braancce-baedd Wasserstein K K-means and its SDP relaxation have much smaller classification error rate on MNIST for the reason that we explained in Example 3 (cf. Lemma 4).
4 Experiments
4.1
Counter-example in Example 3 revisited
Our first experiment is to back up the claim about the failure of centroid-based Wasserstein K- means in Example 3 through simulations. Instead of using point mass measures that may results in instability for computing the barycenters, we use Gaussian distributions with small variance as a smoothed version. We consider K 2, where cluster Gi consists of my many copies of (Pi, P2) pairs and my many μ3, and cluster Gř consists of my many copies of p4. We choose H4 as the following two- two-dimense -io-dimnniionl mixture of Gaussian distributions μı = 0.5 N (a,,,, - 0.5 (at,2, Σx,2) for i = 1, 2, 3, 4. Due to the space limit, detailed simulation setups and parameters are given in Appendix B. From Table 1, we can observe that Wasserstein SDP has achieved exact recovery for all cases while barycenter-based Wasserstein K-means has only around 40% exact recovery rate among all repetitions. In addition, Wasserstein SDP is more stable than distance- based Wasserstein K K-means. Denote A := W ( ) as the squared distance between μ3 and M for k = 1,2, where μ is the barycenter of G Let A, := max max W2(44, Mb) and A := min min W M4) be the maximum within-cluster distance and the minimum between-cluster distance respectively. From Table the Appendix, we can observe that A, < A", from which we can expect Wasserstein SDP to correctly cluster all data points in the Wassertein space. Moreover, Table shows that about 25% times that the distributions (as pa) in Gj satisfy A Ag. implying those μa to be likely assigned to the wrong cluster, which is consistent with
8

---

Example 3. The experiment results also show that any copy of μa is misclassified whenever exact recovery fails for B-WKM, which means the misclassified rate for μa equals (1 y)), where the exact recovery rate for B- 3-WKM shown in Table 1. Table 6 in the appendix further reports the run time comparison, from which we see that distance-basee approaches are more computationally efficient than the barycenter-based one in our settings.
Table 1:
Exact recovery rates and frequency of A₁ A2 for B-WKM among total 50 repetitions in the counter example. W-SDP: Wasserstein SDP, D- WKKM Distance- based Wasserstein K means, B-WKM: Barycenter-baser Wasserstein K-means, total number of distributions.
W-SDP
D-WKM
B-WKM
Frequency of A A2
101
1.00
0.82
0.40
0.32
202
1.00
0.84
0.34
0.26
303
1.00
0.72
0.46
0,20
4.2 Gaussian distributions
and all cluster size equal. We set the centers of each cluster of Gaussians such that all pairwise distances among the barycenters are all equal, i.e., (N V(k)), (0, v(k))) = for all kı, k2 E (1,2,3, 4) with V = max& || [ llop E [4.5, 5.5]. We fix the dimension = 10 and vary the sample size n = 200, 400, 600. And we set the perturbation parameter t = 10 3 on the covariance matrix. The simulation results reported over 100 times in each setting. Fig. 4 shows the misclassification rate (log-scale) versus the squared distance D between centers. We observe that when the distance between centers of clusters are larger than certain threshold (squared distance 1) 10-3 in this case), then Wasserstein SDP can achieve exact recovery for different n, while the misclassification rate for the two Wasserstein K --menss are stably around 10%. And when the distance between centers of clusters are relatively small, the two Wasserstein K-means and SDP behave similarly.
Next, we simulate random Gaussian measures from model (18) with
Error rate for W-SDP, B-WKM D-WKM
10
B-WKM 200
B-WKM 400
B-WKM 600
Error rate (log scale)
1
W-SDP 200
W-SDP 400
W-SDP 600
D-WKM 200
10
D-WKM 400
D-WKM 600
1
0.5
1.5
D
x x
Figure 4: Mis-classification error versus squared distance D from Wasserstein SDP (W-SDP) and barycenter/distance--t baasd Wasserstein K-means (B-WKM and D-WKM) for clustering Gaussians under n e (200, 400, 600}. Due to the log-scale, 10 corresponds to exact recovery.
4.3 Real-data applications
We consider three benchmark image datasets: MNIST, Fashion-MNIST and USPS handwrit- ing digits. Due to complexity issues, we consider subsets of the whole datasets and randomly choose fixed number of images from each clusters based on 10 replicates for each cases. Here we used Sinkhorn divergence to calculate the Wasserstein distance and the Bregman projection with 100 iterations for computing the barycenters, which efficient and stable for non-degenerate case in practice. For both Wasserstein K-means methods, we use the initialization method in analogue to K-means++ for Fuclidean data, i.e. the first cluster barycenter is chosen uni- formly at random of the distributions, after which each subsequent cluster barycenter is

---

chosen from the remaining distributions with probability proportional to its squared Sinkhorn di- vergence from the distribution's is closest existing cluster barycenter. Codes using MATLAB and Python implementing W-SDP and D-WKM are available at: https://github. http://ggtttuu.c.mmteeeeeee com/Yubo02/ Wasserstein---meannee forenns--or--rurreeeee htw.:orhhhccccccoooooooooaaaaaaa
For MNIST dataset, we choose and fix two clusters: G1 containing the number "O" and G. containing the number "5" for two cases. (1) In Case I, we randomly draw 200 number "0" and 100 number of "5" for each repetition. (2) In Case 2. we double the number and randomly draw 400 number "0" and 200 number of "5" instead. For Fashion-MNIST and USPS handwriting digits dataset, we consider K = 3: Gf containing the "T-shirt/top" (or the number "O" for USPS handwriting). G2 containing the "Trouser" (or the number "5" for USPS handwriting) and G3 containing the "Dress" (or the number "7" for USPS handwriting). The cluster sizes are unbalanced where we randomly choose 200, 100 and 100 number from Gi, G2 and G3 respectively. The error rates are shown Table 2. Detailed setups and more results about F1 scores (analogous to error rates) as well as time costs are placed at Appendix A due to space limit.
From Table 2, we can observe that the performances for Wasserstein SDP (W-SDP) and distance-based Wasserstein K-means (D-WKM) are better compared with barycenter-base Wasserstein K-means (B-WKM) for all cases. And the results from Table 3 in Appendix A by using F1 score are consistent with the results from Table 2. The original K K-means method behaves similar as barycenter- based Wasserstein K K-means in some cases and behaves less preferable for cases such as our experiment on USPS handwriting digits. In particular, the visualization of the clustering results for case has been shown in Fig. 3. From this figure we can find that the classification criterion for B-WKM will end up with the closeness to certain shape of "0", which is characterized by certain angle or the degree of stretch. And this will lead to the high misclassification error for barycenter-based or centroid based Wasserstein K -mans.
Table 2: Error rate (SD) for clustering three benchmark datasets: MNIST, Mushion-MNISS and USPS handwriting digits. MNIST, (MNIST2) refers to the results of Case 1 (Case 2) for MNIST dataset.
W-SDP
D-WKM
B-WKM
KM
MNIST
0.235 (0.045)
0.156 (0.057)
0.310 (0.069)
0.295 (0.066)
MNIST,
0.279 (0.050)
0.185 (0.097)
0.324 (0.032)
0.362 (0.033)
Fashion-MNIST
0.082 (0.020)
0.056 (0.014)
0.141 (0.059)
0.138 (0.099)
USPS handwriting
0.206 (0.020)
0.159 (0.061)
0.240 (0.045)
0.284 (0.025)
5 Discussion
In this paper, we observed and analyzed the peculiar behaviors of Wasserstein barycenters and their results in clustering probability distributions. After that, we proposed the restance-based K- means approach (D-WKW) and its semidefinite program relaxation (W-SDP) by showing the exact recovery results for Gaussians theoretically and numerically. For several real benchmark datasets, we showed results where D- D-WKM and W-SDP could outperform barycenter-basd based K K-means approach (B-WKM). And we focused on one unbalanced case of two clusters from MNIST and analyzed its behavior through visualization of misclassification for the barycenter-based Wasserstein K-means. The corresponding time costs for B-WKM, D-WKM and W-SDP for benchmark datasets suggest that the scalability for them and especially our approaches could be serious when sample size is large. One of our future goals would be addressing the corresponding computational complexity issues for real data. Another goal is to find out more conclusive results where our approaches are preferable within out of the realm of unbalanced clusters for real datasets.
Acknowledgments and Disclosure of Funding
Xiaohui Chen was partially supported by NSF CAREER grant DMS-I752614 Yun Yang was partially supported by NSF grant DMS-2210711
10

---

References
Radoslaw Adamezak. tail inequality for suprema of unbounded empirical processes with appli- cations to Markov chains. Electronic Journal of Probability, 13(mone); 1000 1034, 2008. doi:
10.1214/EJRv13-522 URL https: /ttp///doi.oooo// orgg 10. 1014/22141 v13--21 521.
Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Math. Anal., 43 (2):904-922) 924, 2011.
Daniel Aloise, Amit Deshpande, Pierre Hansen, and Preyas Popat. Np-hardness of euclidean sum-of- squares clustering. Machine learning, 75(2):245-248, 2009.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Itera- tive bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2)AA111-AA) 138, 2015.
Rajendra Bhatia, Tanvi Jain, and Yongdo Lim. On the bures-wasserstein distance between positive definite matrices. Expositiones Mathematicae. 37(2):665-111, 2019. ISSN 0723-0869. doi: https: 00....11000... URL https://ww..ciienciiieettt www.ncedireeet.......... wtt.eeeeccttctttooeeeeeeeeeeeeee science/
Jérémie Bigot, Raúl Gouet, Thierry Klein, and Alfredo López. Geodesic PCA in the Wasserstein space by convex PCA. Annales de I'Institul Henri Poincaré, Probabilités et Statistiques, 53(1):1 26, 2017. doi: 10. 10.1214155AAHP7000
URL https: htwe:oiii.ggggg/ org/10. 1214/15-AHHP700
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Com- munications on Pure and Applied Mathematics, 44(4):375-417, 1991. doi: https://doi.org org/10.
1002/cpa.116044440. URL httpp:://omlineelbbraryy wiley. wwtp/ooommaaaacccccooooc.ooooo.a http://a///1..00//.//.. 1002 3160440402.
Samuel Burer and Renato D. C. Monteiro. nonlinear programming algorithm for solving semidefi- nite programs via low-rank factorization. Mathematical Programmmingg 95(2):329-357, 2003. doi:
10110077/101077-002-0335365 URL httpps///oii..ogg///.00/////////
Elsa Cazelles, Vivien Seguy, Jérémie Bigot, Marco Cuturi, and Nicolas Papadakis. Geodesic pea versus log-pea of histograms in the wasserstein space. SIAM Journal Scientific Comput- ing, 40(():0429-BB)) B456, 2018. 17M1143459.
doi: 10.1137/17M1143459, URL https://doi. org/10.1137.
Frédéric Chazal, Clément Levrard, and Martin Royer. Clustering of measures via mean measure quantization. Electronic Journal of Statistics, 15(1):2060 - 2104, 2021. doi: 10.1214/21-EJS1822 URL http::tt::::/ettt //doii 20g/10.100 11t.//1..11111111 1214/21-
Xiaohui Chen and Yun Yang. Cutoff for exact recovery of gaussian mixture models. IEEE Transac-
tions on Information Theory, 676)):4233-4238,, 2021.
Yaqing Chen, Zhenhua Lin, and Hans-Georg Müller. Wasserstein regression. Journal of the American Statistical Association, 010):1-14, 2021. doi: 10. 00.108800161149922211111928
URL
http:://laa..rrg//00000000000000
2021. 1956937. 1921.9999777
Pierre Del Moral and Angele Niclas. A taylor expansion of the square root matrix functional, 2017.
URL https: / /twxiv.ivvvv org/ abb//705 08561.
G. Domazakis, Dimosthenis Drivaliaris, Sotirios Koukoulas, G. I. Papayiannis, Andrianos E. Tsekrekos, and Athanasios N. Yannacopoulos. Clustering measure-valued data with wasserstein barycenters. arXiv: Machine Learning. 2019.
Darina Dvinskikh and Daniil Tiapkin. Improved complexity bounds in wasserstein barycenter problem, 2020. URL https:tt/// pelttss:/aant //arxiv.. .awiiii...g/bbbba org/abs/2010. 046777.77....
Yingjie Fei and Yudong Chen. Hidden integrality of sdp relaxation for sub-gaussian mixture models. 00X.000003.0650, 2018.

---

Aude Genevay, Gabriel Peyre, and Marco Cuturi. Learning generative models with sinkhorn di- vergences. In Amos Storkey and Fernando Perez-Cruz, editors, Proceedings of the Twenty- First International Conference on Artificial Intelligence and Statistics, volume 84 of Pro- ceedings of Machine Learning Research, pages 1608-1617. PMIR, 09-11 Apr 2018. URL https://proceedding mlr. www.mmll.......ee....mmm press/v84/ veeeeeevay////e hmm.mmm..
Christophe Giraud and Nicolas Verzelen. Partial recovery bounds for clustering with the relaxed kmeans. arXiv: 1807007777 2018.
Jan-Christian Hütter and Philippe Rigollet. Minimax estimation of smooth optimal transport maps, 2019. URL
05828.
Hicham Janati, Marco Cuturi, and Alexandre Gramfort. Debiased Sinkhorn barycenters. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 4692-4701. PMLR, 13-18 Jul 2020. URL https: / ptwp//eeinnn..m.r......... sttess//vv//// j html.
Young-Heon Kim and Brendan Pass. Wasserstein barycenters over riemannian manifolds. Ad- vances in Mathematics, 307:640-683, 2017. ISSN 0001-8708. doi: https://doi.or//10.
(0166:.amm.2016.1...22 https://// http///www.ciinccerrreetttt wwm///ciencee/rticlee/pppp 50001180885330443..
Khang Le. Huy Nguyen, Quang M Nguyen, Tung Pham, Hung Bui, and Nhat Ho. On ro- bust optimal transport: Computational complexity and barycenter computation. In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Ad- vances in Neural Information Processing Systems, volume 34, pages 21947-21959, Cur- ran Associates, Inc., 2021. URL https: httte::://proceediiie wtwpeeeeeinnnn..unieiieeeee neurips wwwwwps....paaaeeeeeeeeeeeeeeeee cc/ppaaee//222////1le 5800a77867ee22233444466222011000
Paper pdf.
Stuart Lloyd. Least squares quantization in pem. IEEE Transactions Information Theory, 28: 129-137, 1982.
John Lott. Some geometric calculations on wasserstein space. Communications in Mathematical Physics, 277 437, 2008. doi: 10.1007/500220-07--336--- URL httth::///o:t /ttp:://///o..r/110... 2000:/0000022000000377--
Yu Lu and Harrison Zhou. Statistical and computational guarantees of lloyd's algorithm and its variants. arXiv:1612.02009, 2016.
J.B. MacQueen. Some methods for classification and analysis of multivariate observations. Proc. Fifth Berkeley Sympos. Math. Statist. and Probability, pages 281-297, 1967.
Robert J. McCann. A convexity principle for interacting gases. Advances in Mathematics, 128(1): 153- 1997. ISSN 0001-8708. doi: http:://koo.rrg///000006/aaa1111 1634. URL https:
Marina Meila and Jianbo Shi. Learning segmentation by random walks. In In Advances in Neural Information Processing Systems, pages 873-879. MIT Press, 2001.
Andrew Y. Ng. Michael Jordan, and Yair Weiss. On spectral clustering: Analysis and an algorithm. In Advances in Neural Information Processing Systems, pages 849-856. MIT Press, 2001.
Jiming Peng and Yu Wei. Approximating k-means-type clustering via semidefinite programming. 2007.
SIAM J. OPTIM,
Philippe Rigollet and Jonathan Weed. Uncoupled isotonic regression via minimum Wasserstein deconvolution. Information and Inference: A Journal of the IMA, 8(4),691 04 2019. ISSN 2049-8772. doi: 10.1093/imaiaa/nnz/66 URL. https://doi. org http://..000mmmmamaaaaaaaaaa
Filippo Santambrogio and Xu-Jia Wang. Convexity of the support of the displacement inter-
polation: Counterexamples. Applied Mathematics Letters, 58:152-158, 2016. ISSN 0893- 9659. doi: https://di..org/1..0006.....im
URL https: /wwww..cieecedirec...
com/
12

---

Vivien Seguy and Marco Cuturi. Principal geodesic analysis for probability measures un- der the optimal transport metric. In C. Cortes. N. Lawrence, Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28. Cur- ran Associates, Inc., 2015. URL https::/ ppwpee/e////ceeeg.....e. wwwprrrreednnnne..errpppp.p heeripss.coppper/200///1. /ece/ppper/21//ftlee/ (rcabbbffa13yccllfffee66664ce
pdf.
Justin Solomon, Fermando de Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. ACM Trans. Graph., 34(4), jul 2015. ISSN 0730-0301. doi: 10.1145/2766963. URL httph:tt///// / /doi. org/10 1145/2766963
Aad W. van der Vaart and Jon A. Wellner. Maximal Inequalities and Covering Numbers, pages 95-106. Springer New York, New York, NY, 1996. ISBN 978-1-4757-2545-2. doi: 10. 10.1007/ 978-1-4757-2545-2. 14. URL https://dot org/10 10. 0t11//977---477--258 2_14.
Santosh Vempala and Grant Wang. A spectral algorithm for learning mixture models. J. Comput. Syst. Sci, 68:2004, 2004.
Isabella Verdinelli and Larry Wasserman. Hybrid Wasserstein distance and fast distribution clustering. Electronic Journal of Statistics, 13(2):5088 - 5119, 2019. doi: 10.1214/1--EIS139, URL
http::://doi.org//// 1211111-111/1
Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.
doi: 10. 108231596.
Cédric Villani. Topics in optimal transportation. Graduate studies in mathematics. American mathematical society, 2003.
Ulrike von Luxburg. tutorial on spectral clustering. Statistics and Computing, 17(4):395-416, 2007.
Ulrike Luxburg, Mikhail Belkin, and Olivier Bousquet. Consistency of spectral clustering. Annals of Statistics, 36(2):555-588) 2008.
Yubo Zhuang, Xiaohui Chen, and Yun Yang. Sketch-and-lift: scalable subsampled semidefinite program for k-means clustering. In Gustau Camps-Valll Francisco R. Ruiz, and Isabel Valera, editors, Proceedings The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pages 9214-9246. PMLR, 28-30 Mar 2022. URL https: :ttpps://procee http://ppoeeeee.eeeeee// mlr wwwpppesssssseee http///11122zhn2g2/01 html. wtwpaaansmmr....esee.eeee
13

---

Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on how to answer these questions. For each question, change the default [TODO] to [Yes], [No], , or [N/A] - You are strongly encouraged to include justification your answer, either by referencing the appropriate section of your paper or providing a brief inline description. For example:
Did you include the license to the code and datasets? [Yes] See Section ??.
Did you include the license the code and datasets? [No] The code and the data are proprietary.
Did you include the license to the code and datasets? [N/A]
Please do modify the questions and only use the provided macros for your answers. Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below,
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL)? [N/A]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [N/A]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "xqqr05hDDuh", "reviewer": "Reviewer_k4tW", "summary": "In this paper, the authors provided the pitfalls of Wasserstein K-means via giving the examples and theoretical proof, and proposed a distance-based formulation of K-means in the Wasserstein space.  In addition, the authors conduct a set of experiments on real and simulation data examples to evaluate the proposal. The main contributions are:\n1. The authors analyze the irregularity and non-robustness of barycenter-based formulation of K-means theoretically and empirically.\n2. To escape the pitfalls of the Wasserstein space, the authors generalize the distance-based formulation of K-means to the Wasserstein space.\n", "questions": "1. Please further add the experiments of this paper.\n2. Please include the detailed proof of the pitfalls. \n3. The conclusion section may be missing. Please include this section in the end of this paper.\n", "limitations": "In this paper, the authors only consider the two categories problem, and conduct the experiments on the real-world dataset. The multiple categories problem also should be considered. ", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The proposed problem is worth exploring\n2. The research is meaningful \nWeaknesses:\n1. In this paper, the authors claim that they provide evidence for pitfalls of barycenter-based Wasserstein K-means, but I cannot find the details of proof for the irregularity and non-robustness.\n2. Experiments are not enough. My main concern is the experiments in this paper. The authors conduct a set of experiments on real and simulation data examples to evaluate the proposal. However, the authors only conduct experiment on one real-world dataset and use error rate to evaluate the proposal. Using more datasets, and considering the imbalanced datasets, adding the F1 score can be more convincing. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "vWKvv5lf1hy", "reviewer": "Reviewer_Arph", "summary": "The paper performs studies on three different implementations of K-means clustering for probability distributions. Conventionally, the problem was solved by finding barycenters. The authors point out that such approaches are unsuitable because of two pitfalls in Wasserstein barycenters. In contrast, the alternative implementations with pairwise distances and SDP relaxation can provide better clustering results.\n", "questions": "The proposed distance-based and SDP-based approaches do not include barycenters or means as optimization variables. Can they still be called \"K-means\"?\n\n\nAlthough the proposed alternatives work better in the given examples, they are more expensive for large-scale data sets because they have a quadratic cost to the number of instances. Optimizing over the cluster membership also makes them more difficult to parallelize. All given examples are small and cannot reveal the drawback. The whole MNIST data set has 70,000 images, but the authors picked only a few hundred.\n\nHow did the authors obtain the pairwise distances for MNIST?", "limitations": "In the checklist, the authors claimed that they described the limitations of their work but did not say where. There is no discussion or conclusion section in the paper, and I cannot find any relevant part for the limitations.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\nThe work points out the difference between Wasserstein K-means and conventional K-means. Good examples are provided. Through the paper, readers can realize the difficulty due to the Wasserstein barycenters, and future research efforts should be devoted to non-barycenter-based clustering algorithms.\n\nWeakness:\nMost illustrative examples are small and artificial. Although distance-based and SDP-based approaches work better for such examples, it is hard to say they are the future for the research problem because they are not scalable.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "g6I-7EGHfg9", "reviewer": "Reviewer_Ln53", "summary": "Paper studies the behavior of two formulations of Wasserstein-K-means: i) a centroid based ii) a distance based formulation. It shows that the centroid-based formulation can have some unexpected behaviors and illustrates the phenomenon on several toy examples. A SDP formulation of the distance based formulation is then proposed in order to speed up the computations. Experiments on simulated data are given, together with a simple real-data application on a MNIST clustering scenario. ", "questions": "See above", "limitations": "As far as I can see, there is no potential negative societal impact of the work. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths of the paper :\n- the paper highlight an interesting behavior of the Wasserstein-k-means, advocating for the use of the distance-based formulation rather than the centroid based one (exhibiting some cases in which they lead to different solutions in the Wasserstein space). To my knowledge, this analysis is new. \n- Illustrations are provided, allowing better understanding the behavior in some specific cases.\n\nWeaknesses:\n- The paper advocates the use of the distance-based formulation but the scenarii that are considered seem rather limited. An in-depth analysis of the cases in which the difference appears and matters is lacking. For instance, it seems that it is related to « unbalanced » scenario, that is to say when the different cluster sizes are different. This behavior should be better investigated. \n- the complexity analysis of the algorithm is not provided. \n- The experimental section is rather limited, with one simple illustration on a MNIST clustering (digit 0 vs. digit 5) provided. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "2O-AUBm9g8q", "reviewer": "Reviewer_dPcg", "summary": "The paper uses k-means like algorithm to cluster probability measures using 2- Wasserstein metric. k-means clustering can be formulated in two ways 1) Centroid Based and 2) Distance based . Both formulations are equivalent in the Euclidean space however in the Wasserstein space the same is not true. The authors show that in the Wasserstein space centroid based k-means is much inferior to the distance based version due to irregularity and non-robustness. This is shown using well crafted examples. The authors empirically demonstrate the superiority of the distance based version. In addition the authors also give an SDP relaxation of the distance based k-means in Wasserstein space and show both theoretically and empirically that when the clusters are well separated , the SDP relaxation formulation can recover the clusters exactly with high probability.", "questions": "1)Did you compare the empirical results with Euclidean k-means algorithm as baseline, used on the probability distribution vectors?\n2) Do you try some experiments with distributions other than Gaussian? Are there any results on exact recovery for other probability measures?", "limitations": "see the above sections", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1) The paper appears technically sound. I did not check the proofs but the main theorem and other results appear correct.\n2) Clustering is a very important problem in Machine learning and hence extending a popular algorithm like k-means to another non-Euclidean metric space, backed with a nice theoretical result will be of interest to the community.\n3) The use of examples with accompanied figures to highlight the important results is appreciable and aids understanding.\n\nWeaknesses:\n1) The paper is difficult to follow for a person not having background knowledge of Wasserstein space. For example terms like Alexandrov curvature, optimal transport map etc. are used without giving additional information or definition. I believe adding a bit more background information and details of notations used will go a long way in improving the readability of the paper for broader audience. \n2) Experiments are done on very small data and also only for k values 2 and 4. More experiments with different values of k and larger datasets (may be more clusters from MNIST ) will strengthen the paper.\n3) Related work for clustering of probability distributions is discussed very briefly. It would be good if the authors can cite works on clustering probability distributions in general (if available) and compare and contrast the Wasserstein k-means method.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept