Original: Wasserstein K -means for clustering probability distributions
Revised: Wasserstein K -means for clustering probability distributions

Original: Anonymous Author(s)

Original: AffiliationAddress

Original: Abstract
Revised: Abstract

Original: Clustering is an important exploratory data analysis technique to group objectsbased on their similarity.
Revised: Clustering is an important exploratory data analysis technique to group objects based on their similarity.

Original: The widely used K -means clustering method relieson some notion of distance to partition data into a fewer number of groups.
Revised: The widely used K -means clustering method relies on some notion of distance to partition data into a fewer number of groups.

Original: Inthe Euclidean space, centroid-based and distance-based formulations of the Kmeans are equivalent.
Revised: In the Euclidean space, centroid-based and distance-based formulations of the Kmeans are equivalent.

Original: In modern machine learning applications, data often ariseas probability distributions and a natural generalization to handle measure-valueddata is to use the optimal transport metric.
Revised: 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.

Original: Due to non-negative Alexandrovcurvature of the Wasserstein space, barycenters suffer from regularity and nonrobustness issues.
Revised: Due to non-negative Alexandrov curvature of the Wasserstein space, barycenters suffer from regularity and nonrobustness issues.

Original: The peculiar behaviors of Wasserstein barycenters may make thecentroid-based formulation fail to represent the within-cluster data points, while themore direct distance-based K -means approach and its semidefinite program (SDP)relaxation are capable of recovering the true cluster labels.
Revised: 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.

Original: In the special caseof clustering Gaussian distributions, we show that the SDP relaxed Wasserstein
Revised: In the special case of clustering Gaussian distributions, we show that the SDP relaxed Wasserstein

Original: K -means can achieve exact recovery given the clusters are well-separated underthe 2 -Wasserstein metric.
Revised: K -means can achieve exact recovery given the clusters are well-separated under the 2 -Wasserstein metric.

Original: Our simulation and real data examples also demonstratethat distance-based K -means can achieve better classification performance overthe standard centroid-based K -means for clustering probability distributions andimages.
Revised: 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.

Original:  Introduction
Revised:  Introduction

Original: Clustering is a major tool for unsupervised machine learning problems and exploratory data analysisin statistics.
Revised: Clustering is a major tool for unsupervised machine learning problems and exploratory data analysis in statistics.

Original: Suppose we observe a sample of data points X 1 , . . .
Revised: Suppose we observe a sample of data points X 1 , . . .

Original: , X n taking values in a metricspace ( X , ∥ · ∥ ) .
Revised: , X n taking values in a metric space ( X , ∥ · ∥ ) .

Original: Suppose there exists a clustering structure G ∗ 1 , . . .
Revised: Suppose there exists a clustering structure G ∗ 1 , . . .

Original: , G ∗ K such that each data point
Revised: , G ∗ K such that each data point

Original: X i belongs to exactly one of the unknown cluster G ∗ k .
Revised: X i belongs to exactly one of the unknown cluster G ∗ k .

Original: The goal of clustering analysis is to recoverthe true clusters G ∗ 1 , . . .
Revised: The goal of clustering analysis is to recover the true clusters G ∗ 1 , . . .

Original: , G ∗ K given the input data X 1 , . . .
Revised: , G ∗ K given the input data X 1 , . . .

Original: , X n .
Revised: , X n .

Original: In the Euclidean space X = R p ,the K -means clustering is a widely used method that achieves the empirical success in manyapplications [MacQueen, 1967].
Revised: In the Euclidean space X = R p , the K -means clustering is a widely used method that achieves the empirical success in many applications [MacQueen, 1967].

Original: In modern machine learning and data science problems suchas computer graphics [Solomon et al., 2015], data exhibits complex geometric features and traditionalclustering methods developed for Euclidean data may not be well suited to analyze such data.
Revised: 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.

Original: In this paper, we consider the clustering problem of probability measures µ 1 , . . . , µ n into K groups.
Revised: In this paper, we consider the clustering problem of probability measures µ 1 , . . . , µ n into K groups.

Original: As a motivating example, the MNIST dataset contains images of handwritten digits 0-9.
Revised: As a motivating example, the MNIST dataset contains images of handwritten digits 0-9.

Original: Normalizingthe greyscale images into histograms as probability measures, a common task is to cluster the images.
Revised: Normalizing the greyscale images into histograms as probability measures, a common task is to cluster the images.

Original: One can certainly apply the Euclidean K -means to the vectorized images.
Revised: One can certainly apply the Euclidean K -means to the vectorized images.

Original: However, this wouldlose important geometric information of the two-dimensional data.
Revised: However, this would lose important geometric information of the two-dimensional data.

Original: On the other hand, theory ofoptimal transport [Villani, 2003] provides an appealing framework to model measure-valued data asprobabilities in many statistical tasks [Domazakis et al., 2019, Chen et al., 2021, Bigot et al., 2017,
Revised: On the other hand, theory of optimal transport [Villani, 2003] provides an 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,

Original: Seguy and Cuturi, 2015, Rigollet and Weed, 2019, Hütter and Rigollet, 2019, Cazelles et al., 2018].
Revised: Seguy and Cuturi, 2015, Rigollet and Weed, 2019, Hütter and Rigollet, 2019, Cazelles et al., 2018].

Original: Background on K -means clustering.
Revised: Background on K -means clustering.

Original: Algorithmically, the K -means clustering have two equivalentformulations in the Euclidean space – centroid-based and distance-based – in the sense that theyboth yield the same partition estimate for the true clustering structure.
Revised: Algorithmically, the K -means clustering have two equivalent formulations in the Euclidean space – centroid-based and distance-based – in the sense that they both yield the same partition estimate for the true clustering structure.

Original: Given the Euclidean data
Revised: Given the Euclidean data

Original: X 1 , . . .
Revised: X 1 , . . .

Original: , X n ∈ R p , the centroid-based formulation of standard K -means can be expressed as
Revised: , X n ∈ R p , the centroid-based formulation of standard K -means can be expressed as

Original: [Equation]
Revised: [Equation]

Original: where clusters { G k } Kk =1 are determined by the Voronoi diagram from { β k } Kk =1 , ¯ X k =| G k | − 1 (cid:80) i ∈
Revised: where clusters { G k } Kk =1 are determined by the Voronoi diagram from { β k } Kk =1 , ¯ X k = | G k | − 1 (cid:80) i ∈

Original: G k X i denotes the centroid of cluster G k , (cid:70) denotes the disjoint union and [ n ] ={ 1 , . . . , n } .
Revised: G k X i denotes the centroid of cluster G k , (cid:70) denotes the disjoint union and [ n ] = { 1 , . . . , n } .

Original: Heuristic algorithm for solving (1) includes Lloyd’s algorithm
Revised: Heuristic algorithm for solving (1) includes Lloyd’s algorithm

Original: [Lloyd, 1982], which isan iterative procedure alternating the partition and centroid estimation steps.
Revised: [Lloyd, 1982], which is an iterative procedure alternating the partition and centroid estimation steps.

Original: Specifically, given aninitial centroid estimate β (1)1 , . . .
Revised: Specifically, given an initial centroid estimate β (1)1 , . . .

Original: , β (1) K , one first assigns each data point to its nearest centroid at thet -th iteration according to the Voronoi diagram, i.e.,
Revised: , β (1) K , one first assigns each data point to its nearest centroid at the t -th iteration according to the Voronoi diagram, i.e.,

Original: [Equation]
Revised: [Equation]

Original: and then update the centroid for each cluster
Revised: and then update the centroid for each cluster

Original: [Equation]
Revised: [Equation]

Original: where | G ( t ) k | denotes the cardinality of G ( t ) k .
Revised: where | G ( t ) k | denotes the cardinality of G ( t ) k .

Original: Step (2) and step (3) alternates until convergence.
Revised: Step (2) and step (3) alternate until convergence.

Original: The distance-based (sometimes also referred as partition-based ) formulation directly solves thefollowing constrained optimization problem without referring to the estimated centroids:
Revised: The distance-based (sometimes also referred as partition-based ) formulation directly solves the following constrained optimization problem without referring to the estimated centroids:

Original: [Equation]
Revised: [Equation]

Original: Observe that (1) with nearest centroid assignment and (4) are equivalent for the clustering purposedue to the following identity, which extends the parallelogram law from two points to n points,[Equation] and X i ∈ R
Revised: 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,

Original: [Equation]
Revised: [Equation]

Original: Consequently, the two criteria yield the same partition estimate for G ∗ 1 , . .
Revised: Consequently, the two criteria yield the same partition estimate for G ∗ 1 , . .

Original: , G ∗ K .
Revised: , G ∗ K .

Original: The key identity (5)establishing the equivalence relies on two facts of the Euclidean space: (i) it is a vector space (i.e.,vectors can be averaged in the linear sense); (ii) it is flat (i.e., zero curvature), both of which areunfortunately not true for the Wasserstein space ( P 2 ( R p ) , W 2 ) that endows the space P 2 ( R p ) of allprobability distributions with finite second moments with the 2 -Wasserstein metric W 2 [Ambrosioet{+
 [Ambrosio et+} al., 2005].
Revised: The key identity (5) establishing the equivalence relies on two facts of the Euclidean space: (i) it is a vector space (i.e., vectors can be averaged in the linear sense); (ii) it is flat (i.e., zero curvature), both of which are unfortunately not true for the Wasserstein space ( P 2 ( R p ) , W 2 ) that endows the space P 2 ( R p ) of all probability distributions with finite second moments with the 2 -Wasserstein metric W {+
 [Ambrosio et+} al., 2005].

Original: In particular, the 2-Wasserstein distance between two distributions µ and ν in P 2 ( R p ) isdefined as
Revised: In particular, the 2-Wasserstein distance between two distributions µ and ν in P 2 ( R p ) is defined as

Original: [Equation]
Revised: [Equation]

Original: (6) where minimization over γ runs over all possible couplings with marginals µ and ν .
Revised: where minimization over γ runs over all possible couplings with marginals µ and ν .

Original: It is well-knownthat the Wasserstein space is a metric space (in fact a geodesic space) with non-negative curvature inthe Alexandrov sense [Lott, 2008].
Revised: It is well-known that the Wasserstein space is a metric space (in fact a geodesic space) with non-negative curvature in the Alexandrov sense [Lott, 2008].

Original: Our contributions.
Revised: Our contributions.

Original: We summarize our main contributions as followings: (i) we provide evidence forpitfalls (irregularity and non-robustness) of barycenter-based Wasserstein K -means, both theoreticallyand empirically, and (ii) we generalize the distance-based formulation of K -means to the Wassersteinspace and establish the exact recovery property of its SDP relaxation for clustering Gaussian measuresunder a separateness lower bound in the 2-Wasserstein distance.
Revised: We summarize our main contributions as followings: (i) we provide evidence for pitfalls (irregularity and non-robustness) of barycenter-based Wasserstein K -means, both theoretically and empirically, and (ii) we 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 a separateness lower bound in the 2-Wasserstein distance.

Original: Existing work.
Revised: Existing work.

Original: 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’salgorithm [Lloyd, 1982], spectral methods [von Luxburg, 2007, Meila and Shi, 2001, Ng et al. ,2001], semidefinite programming (SDP) relaxations 
Revised: 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 [Lloyd, 1982], spectral methods [von Luxburg, 2007, Meila and Shi, 2001, Ng et al., 2001], semidefinite programming (SDP) relaxations [Peng and Wei, 2007], non-convex methods via low-rank matrix factorization [Burer and Monteiro, 2003].

Original: [Peng and Wei, 2007], non-convex methods via low-rank matrix factorization

Original: [Burer and Monteiro, 2003].

Original: Theoretic guarantees of those methods areestablished for statistical models on Euclidean data [Lu and Zhou, 2016, von Luxburg et al., 2008,
Revised: Theoretic guarantees of those methods are established for statistical models on Euclidean data [Lu and Zhou, 2016, von Luxburg et al., 2008,

Original: Vempala and Wang, 2004, Fei and Chen, 2018, Giraud and Verzelen, 2018, Chen and Yang, 2021,
Revised: Vempala and Wang, 2004, Fei and Chen, 2018, Giraud and Verzelen, 2018, Chen and Yang, 2021,

Original: Zhuang et al., 2022].
Revised: Zhuang et al., 2022].

Original: To cluster probability measures in the Wasserstein space, the centroid-based Wasserstein K -meansalgorithm has been proposed in Domazakis et al.
Revised: 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.

Original: [2019], which replaced the Euclidean normand sample means by the Wasserstein distance and barycenters respectively.
Revised: It replaced the Euclidean norm and sample means by the Wasserstein distance and barycenters respectively.

Original: More discussions of Wasserstein K -means using barycenters can be found in Section 2.1.
Revised: Verdinelli and

Original: Verdinelli and Wasserman[2019] proposed a modified Wasserstein distance for distribution clustering. 
Revised: Wasserman [2019] proposed a modified Wasserstein distance for distribution clustering. And after that, Chazal et al.

Revised: [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.

Revised: 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.

Original:  Wasserstein K -means clustering methods
Revised:  Wasserstein K -means clustering methods

Original: In this section, we generalize the Euclidean K -means to the Wasserstein space.
Revised: In this section, we generalize the Euclidean K -means to the Wasserstein space.

Original: Our starting point isto mimic the standard K -means methods for Euclidean data.
Revised: Our starting point is to mimic the standard K -means methods for Euclidean data.

Original: Thus we may define two versions of the
Revised: Thus we may define two versions of the

Original: Wasserstein K -means clustering formulations: centroid-based and distance-based .
Revised: Wasserstein K -means clustering formulations: centroid-based and distance-based .

Original: As we mentionedin Section 1, when working with Wasserstein space ( P 2 ( R p ) , W 2 ) , the corresponding centroid-basedcriterion (1) and the distance-based criterion (4), where the Euclidean metric ∥ · ∥ 2 is replaced withthe 2 -Wasserstein metric W 2 , may lead to radically different clustering schemes.
Revised: As we mentioned in Section 1, when working with Wasserstein space ( P 2 ( R p ) , W 2 ) , the corresponding centroid-based criterion (1) and the distance-based criterion (4), where the Euclidean metric ∥ · ∥ 2 is replaced with the 2 -Wasserstein metric W 2 , may lead to radically different clustering schemes.

Original: To begin with, wewould like to argue that due to the irregularity and non-robustness of barycenters in the Wassersteinspace, the centroid-based criterion may lead to unreasonable clustering schemes that lack physicalinterpretations and are sensitive to small data perturbations.
Revised: To begin with, we would like to argue that due to the irregularity and non-robustness 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.

Original: 2.1 Clustering based on barycenters
Revised: 2.1 Clustering based on barycenters

Original: The centroid-based Wasserstein K -means for extending the Lloyd’s algorithm into the Wassersteinspace has been recently considered by Domazakis et al.
Revised: The centroid-based Wasserstein K -means for extending the Lloyd’s algorithm into the Wasserstein space has been recently considered by Domazakis et al.

Original: Specifically, it is an iterativealgorithm proceeds as following.
Revised: Specifically, it is an iterative algorithm proceeds as following.

Original: Given an initial centroid estimate ν (1)1 , . . .
Revised: Given an initial centroid estimate ν (1)1 , . . .

Original: , ν (1) K , one first assignseach probability measure µ 1 , . .. , µ n to its nearest centroid in the Wasserstein geometry at the t -thiteration according to the Voronoi diagram:
Revised: , ν (1) K , one first assigns each probability measure µ 1 , . . µ n to its nearest centroid in the Wasserstein geometry at the t -th iteration according to the Voronoi diagram:

Original: [Equation]
Revised: [Equation]

Original: and then update the centroid for each cluster
Revised: and then update the centroid for each cluster

Original: [Equation]
Revised: [Equation]

Original: Note that ν ( t +1) k in (8) is referred as barycenter of probability measures µ i , i ∈ G ( t ) k , a Wassersteinanalog of the Euclidean average or mean [Agueh and Carlier, 2011].
Revised: Note that ν ( t +1) k in (8) is referred as barycenter of probability measures µ i , i ∈ G ( t ) k , a Wasserstein analog of the Euclidean average or mean [Agueh and Carlier, 2011].

Original: We will also ex-changeablyuse barycenter-based K -means to mean the centroid-based K-means in the Wasserstein space.
Revised: We will also ex-changeably use barycenter-based K -means to mean the centroid-based K-means in the Wasserstein space.

Original: Eventhough the Wasserstein barycenter is a natural notion of averaging probability measures, it mayexhibit peculiar behaviours and fail to represent the within-cluster data points, partly due to theviolation of the generalized parallelogram law (5) (for non-flat spaces) if the Euclidean metric ∥ · ∥ 2is replaced with the 2 -Wasserstein metric W 2 .
Revised: Even though the Wasserstein barycenter is a 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) if the Euclidean metric ∥ · ∥is replaced with the 2 -Wasserstein metric W 2 .

Original: Example 1 ( Irregularity of Wasserstein barycenters) .
Revised: Example 1 ( Irregularity of Wasserstein barycenters) .

Original: Wasserstein barycenter has much less regularity than the sample mean in the Euclidean space
Revised: Wasserstein barycenter has much less regularity than the sample mean in the Euclidean space

Original: [Kim and Pass, 2017].
Revised: [Kim and Pass, 2017].

Original: In particular, Santambrogioand Wang [2016] constructed a simple example of two probability measures that are supported on linesegments in R 2 , whereas the support of their barycenter obtained as the displacement interpolationthe two endpoint probability measures is not convex (cf.
Revised: In particular, Santambrogio and Wang [2016] constructed a simple example of two probability measures that are supported on line segments in R 2 , whereas the support of their barycenter obtained as the displacement interpolation the two endpoint probability measures is not convex (cf.

Original: left plot in Figure 1).
Revised: left plot in Figure 1).

Original: In this example, theprobability density µ 0 and µ 1 are supported on the line segments L 0 = { ( s, as ) : s ∈ [ − 1 , 1] } and
Revised: In this example, the probability density µ 0 and µ 1 are supported on the line segments L 0 = { ( s, as ) : s ∈ [ − 1 , 1] } and

Original: L 1 = { ( s, − as ) : s ∈ [ − 1 , 1] } respectively.
Revised: L 1 = { ( s, − as ) : s ∈ [ − 1 , 1] } respectively.

Original: We choose a ∈ (0 , 1) to identify the orientation of L 0and L 1 based on the x -axis.
Revised: We choose a ∈ (0 , 1) to identify the orientation of Land L 1 based on the x -axis.

Original: Moreover, we consider the linear density functions µ 0 ( s ) =
Revised: Moreover, we consider the linear density functions µ 0 ( s ) =

Original: (1 − s ) / 2and µ 1 ( s ) =
Revised: (1 − s ) / and µ 1 ( s ) =

Original: (1 + s ) / 2 for s ∈ [ − 1 , 1] supported on L 0 and L 1 respectively.
Revised: (1 + s ) / 2 for s ∈ [ − 1 , 1] supported on L 0 and L 1 respectively.

Original: Then the optimaltransport map T :
Revised: Then the optimal transport map T :

Original: = T µ 0 → µ 1 from µ 0 to µ 1 is given by
Revised: = T µ 0 → µ 1 from µ 0 to µ 1 is given by

Original: [Equation]
Revised: [Equation]

Original: and the barycenter corresponds to the displacement interpolation µ t =
Revised: and the barycenter corresponds to the displacement interpolation µ t =

Original: [(1 − t ) id + tT ]
Revised: [(1 − t ) id + tT ]

Original: ♯ µ 0 att = 0 .
Revised: ♯ µ 0 at t = 0 .

Original: [McCann, 1997].
Revised: [McCann, 1997].

Original: For self-contained purpose, we give the proof of (9) in Appendix D.1.
Revised: For self-contained purpose, we give the proof of (9) in Appendix D.1.

Original: Fig.
Revised: Fig.

Original: 1 on the left shows the support of barycenter µ 0 .
Revised: 1 on the left shows the support of barycenter µ 0 .

Original: 5 is not convex (in fact part of an ellipseboundary) even though the supports of µ 0 and µ 1 are convex.
Revised: 5 is not convex (in fact part of an ellipse boundary) even though the supports of µ 0 and µ 1 are convex.

Original: This example shows that the barycenterfunctional is not geodesically convex in the Wasserstein space.
Revised: This example shows that the barycenter functional is not geodesically convex in the Wasserstein space.

Original: As barycenters turn out to be essentialin centroid-based Wasserstein K -means and irregularity of the barycenter may fail to represent thecluster (see more details in Example 3 and Remark 9 below), this counter-example is our motivationto seek alternative formulation. 
Revised: As barycenters turn out to be essential in centroid-based Wasserstein K -means and irregularity of the barycenter may fail to represent the cluster (see more details in Example 3 and Remark 9 below), this counter-example is our motivation to seek alternative formulation. ■

Original: [Figure]
Revised: [Figure]

Revised: 5 m 1 , a =

Original: Example 2 ( Non-robustness of Wasserstein barycenters) .
Revised: Example 2 ( Non-robustness of Wasserstein barycenters) .

Original: Another unappealing feature of the
Revised: Another unappealing feature of the

Original: Wasserstein barycenter is its sensitivity to data perturbation: a small (local) change in one contributingprobability measure may lead to large (global) changes in the resulting barycenter.
Revised: 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.

Original: See Fig.
Revised: See Fig.

Original: 1 onthe right for such an example.
Revised: 1 on the right for such an example.

Original: In this example, we take the source measure as µ 0 = 0 .
Revised: In this example, we take the source measure as µ 0 = 0 .

Original: 5 δ ( − 1 , 1) +0 .
Revised: 5 δ ( − 1 , 1) + 0 .

Original: 5 δ (1 , − 1) and the target measure as µ 1 = 0 .
Revised: 5 δ (1 , − 1) and the target measure as µ 1 = 0 .

Original: 5 δ ( − 1 , − a ) + 0 .
Revised: 5 δ ( − 1 , − a ) + 0 .

Original: 5 δ (1 ,a ) for some a > 0 .
Revised: 5 δ (1 ,a ) for some a > 0 .

Original: It is easy to seethat the optimal transport map T :
Revised: It is easy to see that the optimal transport map T :

Original: = T µ 0 → µ 1 has a dichotomy behavior: 
Revised: = T µ 0 → µ 1 has a dichotomy behavior: [Equation] a

Original: [Equation]
Revised: [Equation]

Revised: and T (1 , −

Revised: [Equation]

Revised: [Equation]

Revised: a

Revised: [Equation]

Original:  Thus the Wasserstein barycenter determined by the displacement interpolation µ t = [(1 − t ) id + tT ]
Revised: [Equation] a [Equation] Thus the Wasserstein barycenter determined by the displacement interpolation µ t = [(1 − t ) id + tT ]

Original: ♯ µ 0is a discontinuous function at a = 1 .
Revised: ♯ µis a discontinuous function at a = 1 .

Original: This non-robustness can be attribute to the discontinuity ofthe Wasserstein barycenter as a function of its input probability measures; in contrast, the Euclideanmean is a globally Lipchitz continuous function of its input points.
Revised: This non-robustness can be attribute to the discontinuity of the Wasserstein barycenter as a function of its input probability measures; in contrast, the Euclidean mean is a globally Lipchitz continuous function of its input points.

Original: ■
Revised: ■

Original: Because of these pitfalls of the Wasserstein barycenter shown in Examples 1 and 2, the centroidbased Wasserstein K -means approach described at the beginning of this subsection may lead tounreasonable and unstable clustering schemes.
Revised: Because of these pitfalls of the Wasserstein barycenter shown in Examples 1 and 2, the centroidbased Wasserstein K -means approach described at the beginning of this subsection may lead to unreasonable and unstable clustering schemes.

Original: In addition, an ill-conditioned configuration maysignificantly slow down the convergence of commonly used barycenter approximating algorithmssuch as iterative Bregman projections [Benamou et al., 2015].
Revised: In addition, an ill-conditioned configuration may significantly slow down the convergence of commonly used barycenter approximating algorithms such as iterative Bregman projections [Benamou et al., 2015].

Original: Below, we give a concrete example ofsuch phenomenon in the clustering context.
Revised: Below, we give a concrete example of such phenomenon in the clustering context.

Original: Example 3 ( Failure of centroid-based Wasserstein K -means) .
Revised: Example 3 ( Failure of centroid-based Wasserstein K -means) .

Original: In a nutshell, the failure in thisexample is due to the counter-intuitive phenomenon illustrated in the right panel of Fig.
Revised: In a nutshell, the failure in this example is due to the counter-intuitive phenomenon illustrated in the right panel of Fig.

Original: 2, wheresome distribution µ 3 in the Wasserstein space may have larger W 2 distance to Wasserstein barycenterµ ∗ 1 than every distribution µ i ( i = 1 , 2 ) that together forms it.
Revised: 2, where some distribution µ 3 in the Wasserstein space may have larger W 2 distance to Wasserstein barycenter µ ∗ 1 than every distribution µ i ( i = 1 , 2 ) that together forms it.

Original: As a result of this strange configuration,even though µ 3 is closer to µ 1 and µ 2 from the first cluster with barycenter µ ∗ 1 than µ 4 coming froma second cluster with barycenter µ ∗ 2 , it will be incorrectly assigned to the second cluster using thecentroid-based criterion (7), since W 2 ( µ 3 , µ ∗ 1 ) > W 2 ( µ 3 , µ ∗ 2 ) > max (cid:8) W 2 ( µ 3 , µ 1 ) , W 2 ( µ 3 , µ 2 ) (cid:9) .
Revised: As a result of this strange configuration, even though µ 3 is closer to µ 1 and µ 2 from the first cluster with barycenter µ ∗ 1 than µ 4 coming from a second cluster with barycenter µ ∗ 2 , it will be incorrectly assigned to the second cluster using the centroid-based criterion (7), since W 2 ( µ 3 , µ ∗ 1 ) > W 2 ( µ 3 , µ ∗ 2 ) > max (cid:8) W 2 ( µ 3 , µ 1 ) , W 2 ( µ 3 , µ 2 ) (cid:9) .

Original: In contrast, for Euclidean spaces due to the following equivalent formulation of the generalizedparallelogram law (5),[Equation] for any X
Revised: In contrast, for Euclidean spaces due to the following equivalent formulation of the generalized parallelogram law (5),

Original: [Equation]
Revised: [Equation]

Original:  R [Equation] there is always some point
Revised: for any X ∈ R p , there is always some point

Original: X i † satisfying ∥
Revised: X i † satisfying ∥

Original: X − X i † ∥ 2 ⩾ ∥
Revised: X − X i † ∥ 2 ⩾ ∥

Original: X − ¯ X ∥ 2 , that is, further away from Xthan the mean ¯ X ; thereby excluding counter-intuitive phenomena as the one shown in Fig.
Revised: X − ¯ X ∥ 2 , that is, further away from X than the mean ¯ X ; thereby excluding counter-intuitive phenomena as the one shown in Fig.

Original: Concretely, the first cluster G ∗ 1 is shown in the left panel of Fig.
Revised: Concretely, the first cluster G ∗ 1 is shown in the left panel of Fig.

Original: 2 highlighted by a red circle,consisting of m copies of ( µ 1 , µ 2 ) pairs and one copy of µ 3) ; the second cluster G ∗ 2 containingcopies of µ 4 is highlighted by a blue circle.
Revised: 2 highlighted by a red circle, consisting of m copies of ( µ 1 , µ 2 ) pairs and one copy of µ 3 ; the second cluster G ∗ 2 containing copies of µ 4 is highlighted by a blue circle.

Original: Each distribution assigns equal probability mass to twopoints, where the two supporting points are connected by a dashed line for easy illustration.
Revised: Each distribution assigns equal probability mass to two points, where the two supporting points are connected by a dashed line for easy illustration.

Original: Morespecifically, we set
Revised: More specifically, we set

Original: [Equation]
Revised: [Equation]

Original: where δ ( x,y ) denotes the point mass measure at point ( x, y ) , and ( x, y, ϵ 1 , ϵ 2 ) are positive constants.
Revised: where δ ( x,y ) denotes the point mass measure at point ( x, y ) , and ( x, y, ϵ 1 , ϵ 2 ) are positive constants.

Original: The property of this configuration can be summarized by the following lemma.
Revised: The property of this configuration can be summarized by the following lemma.

Original: Lemma 4 ( Configuration characterization) .
Revised: Lemma 4 ( Configuration characterization) .

Original: If ( x, y, ϵ 1 , ϵ 2 ) satisfiesy [Equation] < min { x 2 , 0 .
Revised: If ( x, y, ϵ 1 , ϵ 2 ) satisfies y 2 < min { x 2 , 0 .

Original:  ∆ ϵ 1 ,x } and ∆ ϵ 1 ,x < ϵ 22 < ∆ ϵ 1 ,x + y 2 , where ∆ ϵ 1 ,x :=
Revised: 25 ∆ ϵ 1 ,x } and ∆ ϵ 1 ,x < ϵ 22 < ∆ ϵ 1 ,x + y 2 , where ∆ ϵ 1 ,x :=

Original: ϵ 21 + 2 x 2 + 2 xϵ 1 , then for all sufficiently large m (number of copies of µ 1 and µ 2 ),
Revised: ϵ 21 + 2 x 2 + 2 xϵ 1 , then for all sufficiently large m (number of copies of µ 1 and µ 2 ),

Original: W 2 ( µ 3 , µ ∗ 2 ) <
Revised: W 2 ( µ 3 , µ ∗ 2 ) <

Original: W 2 ( µ 3 , µ ∗ 1 ) and
Revised: W 2 ( µ 3 , µ ∗ 1 ) and

Original: [Equation]
Revised: [Equation]

Original: largest 
Revised: largest within

Original: [Equation]
Revised: [Equation]

Revised: least between

Original: [Table]
Revised: [Equation]

Original:  where µ ∗ k denotes the Wasserstein barycenter of cluster G k for k = 1 , 2 .
Revised: cluster [Equation] where µ ∗ k denotes the Wasserstein barycenter of cluster G k for k = 1 , 2 .

Original: Note that the condition of Lemma 4 implies y < x .
Revised: Note that the condition of Lemma 4 implies y < x .

Original: Therefore, the barycenter between µ 1 and µ 2is ˜ µ ∗ 1 := 0 .
Revised: Therefore, the barycenter between µ 1 and µis ˜ µ ∗ 1 := 0 .

Original: 5 δ ( x, 0) + 0 .
Revised: 5 δ ( x, 0) + 0 .

Original: 5 δ ( − x, 0) lying on the horizontal axis.
Revised: 5 δ ( − x, 0) lying on the horizontal axis.

Original: By increasing m , the barycenter µ ∗ 1of cluster G ∗ 1 can be made arbitrarily close to ˜ µ ∗ .
Revised: By increasing m , the barycenter µ ∗of cluster G ∗ 1 can be made arbitrarily close to ˜ µ ∗ .

Original: The second inequality in Lemma 4 shows thatall within-cluster distances are strictly less than the between-cluster distances; therefore, clusteringbased on pairwise distances is able to correctly recover the cluster label of µ 3 .
Revised: The second inequality in Lemma 4 shows that all within-cluster distances are strictly less than the between-cluster distances; therefore, clustering based on pairwise distances is able to correctly recover the cluster label of µ 3 .

Original: However, since µ 3is closer to the barycenter µ ∗ 2 of cluster G ∗ 2 according to the first inequality in Lemma 4, it willbe mis-classified into G ∗ 2 using the centroid-based criterion.
Revised: However, since µis closer to the barycenter µ ∗ 2 of cluster G ∗ 2 according to the first inequality in Lemma 4, it will be mis-classified into G ∗ 2 using the centroid-based criterion.

Original: We emphasize that cluster positionsin this example are generic and do exist in real data; see Remark 9 and Section 4.3 for furtherdiscussions on our experiment results on MNIST data.
Revised: 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.

Original: Moreover, similar to Example 2, a smallchange in the orientation of distribution µ 1 may completely alter the clustering membership of µ 3based on the centroid criterion.
Revised: Moreover, similar to Example 2, a small change in the orientation of distribution µ 1 may completely alter the clustering membership of µbased on the centroid criterion.

Original: Specifically, if we slightly increase x to make it exceed y , thenthe barycenter between µ 1 and µ 2 becomes ¯ µ ∗ 1 := 0 .
Revised: Specifically, if we slightly increase x to make it exceed y , then the barycenter between µ 1 and µ 2 becomes ¯ µ ∗ 1 := 0 .

Original: 5 δ (0 ,y ) + 0 .
Revised: 5 δ (0 ,y ) + 0 .

Original: 5 δ (0 , − y ) that lies on the verticalaxis.
Revised: 5 δ (0 , − y ) that lies on the vertical axis.

Original: Correspondingly, if based on centroids, then µ 3 should be clustered into G ∗ 1 as it is closer tothe barycenter µ ∗ 1 of G ∗ 1 than the barycenter µ ∗ 2 of G ∗ 2 .
Revised: Correspondingly, if based on centroids, then µ 3 should be clustered into G ∗ 1 as it is closer to the barycenter µ ∗ 1 of G ∗ 1 than the barycenter µ ∗ 2 of G ∗ 2 .

Original: Therefore, the centroid-based criterion canbe unstable against data perturbations.
Revised: Therefore, the centroid-based criterion can be unstable against data perturbations.

Original: In comparison, a pairwise distances based criterion alwaysassigns µ 3 into cluster G ∗ 2 no matter x
Revised: In comparison, a pairwise distances based criterion always assigns µ 3 into cluster G ∗ 2 no matter x

Original: < y or x > y .
Revised: < y or x > y .

Original: [Figure] [Equation] [Figure]
Revised: ■

Original:  based on pairwise distances
Revised: 2.2 Clustering based on pairwise distances

Original: Due to the irregularity and non-robustness of centroid-based Wasserstein K -means, we insteadpropose and advocate the use of distance-based Wasserstein K -means below, which extends the
Revised: Due to the irregularity and non-robustness of centroid-based Wasserstein K -means, we instead propose and advocate the use of distance-based Wasserstein K -means below, which extends the

Original: Euclidean distance-based K -means formulation (4) into the Wasserstein space,
Revised: Euclidean distance-based K -means formulation (4) into the Wasserstein space,

Original: [Equation]
Revised: [Equation]

Original: Correspondingly, we can analogously design a greedy algorithm resembling the Wasserstein Lloyd’salgorithm described in Section 2.1 that solves the centroid-based Wasserstein K -means.
Revised: Correspondingly, we can analogously design a greedy algorithm resembling the Wasserstein Lloyd’s algorithm described in Section 2.1 that solves the centroid-based Wasserstein K -means.

Original: Specifically,the greedy algorithm proceeds in an iterative manner as following.
Revised: Specifically, the greedy algorithm proceeds in an iterative manner as following.

Original: Given an initial cluster membershipestimate G (1)1 , . . .
Revised: Given an initial cluster membership estimate G (1)1 , . . .

Original: , G (1) K , one assigns each probability measure µ 1 , . .
Revised: , G (1) K , one assigns each probability measure µ 1 , . .

Original: . , µ n based on minimizing theaveraged squared W 2 distances to all current members in every cluster, leading to an updated clustermembership estimate
Revised: . , µ n based on minimizing the averaged squared W 2 distances to all current members in every cluster, leading to an updated cluster membership estimate

Original: [Equation]
Revised: [Equation]

Original: We arbitrarily select among the least W 2 distance clusters in the case of a tie.
Revised: We arbitrarily select among the least W 2 distance clusters in the case of a tie.

Original: We highlight thatthe center-based and distance-based Wasserstein K -means formulations may not necessarily beequivalent to yield the same cluster labels (cf.
Revised: 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.

Original: Example 3).
Revised: Example 3).

Original: Below, we shall give some exampleillustrating connections to the standard K -means clustering in the Euclidean space.
Revised: Below, we shall give some example illustrating connections to the standard K -means clustering in the Euclidean space.

Original: Example 5 ( Degenerate probability measures) .
Revised: Example 5 ( Degenerate probability measures) .

Original: If the probability measures are Dirac at point
Revised: If the probability measures are Dirac at point

Original: X i ∈ R p , i.e., µ i = δ X
Revised: X i ∈ R p , i.e., µ i = δ X

Original: i , then the Wasserstein K -means is the same as the standard K -means since
Revised: i , then the Wasserstein K -means is the same as the standard K -means since

Original: W 2 ( µ i , µ j ) =
Revised: W 2 ( µ i , µ j ) =

Original: ∥
Revised: ∥

Original: X i − X j ∥ 2 .
Revised: X i − X j ∥ 2 .

Original: ■
Revised: ■

Original: Example 6 ( Gaussian measures) .
Revised: Example 6 ( Gaussian measures) .

Original: If µ i = N ( m i , V i ) with positive-definite covariance matricesΣ i ≻ 0 , then the squared 2 -Wasserstein distance can be expressed as the sum of the squared Euclideandistance on the mean vector and[Equation]
Revised: If µ i = N ( m i , V i ) with positive-definite covariance matrices Σ i ≻ 0 , then the squared 2 -Wasserstein distance can be expressed as the sum of the squared Euclidean distance on the mean vector and

Original: [Equation]
Revised: [Equation]

Original: the squared Bures distance on the covariance matrix [Bhatia et al., 2019].
Revised: the squared Bures distance on the covariance matrix [Bhatia et al., 2019].

Original: Here, we use V 1 / 2 todenote the unique symmetric square root matrix of V ≻ 0 .
Revised: Here, we use V 1 / 2 to denote the unique symmetric square root matrix of V ≻ 0 .

Original: That is,
Revised: That is,

Original: [Equation]
Revised: [Equation]

Original: [Equation] Then the Wasserstein K -means, formulated either in (7) or (11), can be viewed as a covarianceadjusted Euclidean K -means by taking account into the shape or orientation information in the(non-degenerate) Gaussian inputs.
Revised: Then the Wasserstein K -means, formulated either in (7) or (11), can be viewed as a covarianceadjusted Euclidean K -means by taking account into the shape or orientation information in the (non-degenerate) Gaussian inputs.

Original: ■
Revised: ■

Original: Example 7 ( One-dimensional probability measures) .
Revised: Example 7 ( One-dimensional probability measures) .

Original: If µ i are probability measures on R withcumulative distribution function (cdf) F i , then the Wasserstein distance can be written in terms of thequantile transform
Revised: If µ i are probability measures on R with cumulative distribution function (cdf) F i , then the Wasserstein distance can be written in terms of the quantile transform

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where F − is the generalized inverse of the cdf F on [0 , 1] defined as F − ( u )
Revised: where F − is the generalized inverse of the cdf F on [0 , 1] defined as F − ( u )

Original: = inf { x ∈ R : F ( x ) >u } (cf.
Revised: = inf { x ∈ R : F ( x ) > u } (cf.

Original: Theorem 2.18 [Villani, 2003]).
Revised: Theorem 2.18 [Villani, 2003]).

Original: Thus the one-dimensional probability measures in Wassersteinspace can be isometrically embedded in a flat L 2 space, and we can bring back the equivalence of the
Revised: Thus the one-dimensional probability measures in Wasserstein space can be isometrically embedded in a flat L 2 space, and we can bring back the equivalence of the

Original: Wasserstein and Euclidean K -means clustering methods.
Revised: Wasserstein and Euclidean K -means clustering methods.

Original: ■
Revised: ■

Original:  SDP relaxation and its theoretic guarantee
Revised:  SDP relaxation and its theoretic guarantee

Original: Note that Wasserstein Lloyd’s algorithm requires to use and compute the barycenter in (7) and (8)at each iteration, which can be computationally expensive when the domain dimension d is largeor the configuration is ill-conditioned (cf.
Revised: Note that Wasserstein Lloyd’s algorithm requires to use and compute the barycenter in (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.

Original: Example 2).
Revised: Example 2).

Original: On the other hand, it is known that solvingthe distance-based K -means (4) is worst-case NP-hard for Euclidean data.
Revised: On the other hand, it is known that solving the distance-based K -means (4) is worst-case NP-hard for Euclidean data.

Original: Thus we expect solvingthe distance-based Wasserstein K -means (11) is also computationally hard.
Revised: Thus we expect solving the distance-based Wasserstein K -means (11) is also computationally hard.

Original: A common way isto consider convex relaxations to approximate the solution of (11).
Revised: A common way is to consider convex relaxations to approximate the solution of (11).

Original: It is known that certain SDPrelaxation is information-theoretically tight for (4) when the data X 1 , . . .
Revised: It is known that certain SDP relaxation is information-theoretically tight for (4) when the data X 1 , . . .

Original: , X n ∈ R p are generatedfrom a Gaussian mixture model with isotropic known variance [Chen and Yang, 2021].
Revised: , X n ∈ R p are generated from a Gaussian mixture model with isotropic known variance [Chen and Yang, 2021].

Original: In this paper,we extend the idea into Wasserstein setting for solving (11).
Revised: In this paper, we extend the idea into Wasserstein setting for solving (11).

Original: A typical SDP relaxation for Euclidean data uses pairwise inner products to construct an affinity matrixfor clustering [Peng and Wei, 2007]; unfortunately, due to the non-flatness nature, a globally welldefined inner product does not exist for Wasserstein spaces with dimension higher than one.
Revised: A 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 welldefined inner product does not exist for Wasserstein spaces with dimension higher than one.

Original: Therefore,we will derive a Wasserstein SDP relaxation to the combinatorial optimization problem (4) using the squared distance matrix
Revised: Therefore, we will derive a Wasserstein SDP relaxation to the combinatorial optimization problem (4) using the squared distance matrix

Original: A n × n = { a ij } with a ij =
Revised: A n × n = { a ij } with a ij =

Original: W 22 ( µ i , µ j ) .
Revised: W 22 ( µ i , µ j ) .

Original: Concretely, we can one-to-onereparameterize any partition ( G 1 , . . .
Revised: Concretely, we can one-to-one reparameterize any partition ( G 1 , . . .

Original: , G K ) as a binary assignment matrix H = { h ik } ∈ { 0 , 1 } n × Ksuch that h
Revised: , G K ) as a binary assignment matrix H = { h ik } ∈ { 0 , 1 } n × K

Original:  ik = 1 if i ∈ G k and h ik = 0 otherwise.
Revised: such that h ik = 1 if i ∈ G k and h ik = 0 otherwise.

Original: Then (11) can be expressed as a nonlinear 0-1integer program,
Revised: Then (11) can be expressed as a nonlinear 0integer program,

Original: [Equation]
Revised: [Equation]

Original:  where 1 n is the n × 1 vector of all ones and B = diag ( | G 1 | − 1 , . .
Revised: (16) where 1 n is the n × 1 vector of all ones and B = diag ( | G 1 | − 1 , . .

Original: . , | G K | − 1 ) .
Revised: . , | G K | − 1 ) .

Original: Changing of variableto the membership matrix Z = HBH ⊤ , we note that Z n × n is a symmetric positive semidefinite(psd)
Revised: Changing of variable to the membership matrix Z = HBH ⊤ , we note that Z n × n is a symmetric positive semidefinite (psd)

Original: matrix Z ⪰ 0 such that Tr ( Z ) =
Revised: matrix Z ⪰ 0 such that Tr ( Z ) =

Original: K, Z 1 n = 1 n , and Z ⩾ 0 entrywise.
Revised: K, Z 1 n = 1 n , and Z ⩾ 0 entrywise.

Original: Thus we obtain the
Revised: Thus we obtain the

Original: SDP relaxation of (11) by only preserving these convex constraints:
Revised: SDP relaxation of (11) by only preserving these convex constraints:

Original: [Equation]
Revised: [Equation]

Original: To theoretically justify the SDP formulation (17) of Wasserstein K -means, we consider the scenarioof clustering Gaussian distributions in Example 6, where the Wasserstein distance (14) containstwo separate components: the Euclidean distance on mean vector and the Bures distance (13)on covariance matrix.
Revised: 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 on mean vector and the Bures distance (13) on covariance matrix.

Original: Without loss of generality, we focus on mean-zero Gaussian distributionssince optimal separation conditions for exact recovery based on the Euclidean mean componenthave been established in [Chen and Yang, 2021].
Revised: 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].

Original: Suppose we observe Gaussian distributionsν i ∼ N (0 , V i ) , i ∈ [ n ] from K groups G ∗ 1 , · · · , G ∗ K , where cluster G ∗ k contains n k members, andthe covariance matrices have the following clustering structure: if i ∈ G ∗ k , then
Revised: Suppose we observe Gaussian distributions ν i ∼ N (0 , V i ) , i ∈ [ n ] from K groups G ∗ 1 , · · · , G ∗ K , where cluster G ∗ k contains n k members, and the covariance matrices have the following clustering structure: if i ∈ G ∗ k , then

Original: V i = ( I + tX i ) V ( k ) ( I + tX i ) with X 1 , . . . , X n i.i.d.
Revised: [Equation]

Original: ∼ SymN (0 , 1) , (18) where the psd matrix V ( k ) is the center of the k -th cluster, SymN (0 , 1) denotes the symmetric randommatrix with i.i.d.
Revised: where the psd matrix V ( k ) is the center of the k -th cluster, SymN (0 , 1) denotes the symmetric random matrix with i.i.d.

Original: standard normal entries, and t is a small perturbation parameter such that
Revised: standard normal entries, and t is a small perturbation parameter such that

Original: ( I + tX i ) ispsd with high probability.
Revised: ( I + tX i ) is psd with high probability.

Original: For zero-mean Gaussian distributions, we have W 2 ( N (0 , V ) , N (0 , U ))
Revised: For zero-mean Gaussian distributions, we have W 2 ( N (0 , V ) , N (0 , U ))

Original: =
Revised: =

Original: d ( V, U ) according to (14).
Revised: d ( V, U ) according to (14).

Original: Note that on the Riemannian manifold of psd matrices, the geodesicemanating from V ( k ) in the direction X as a symmetric matrix can be linearized by V
Revised: Note that on the Riemannian manifold of psd matrices, the geodesic emanating from V ( k ) in the direction X as a symmetric matrix can be linearized by V

Original: = ( I +tX ) V ( k )
Revised: = ( I + tX ) V ( k )

Original: ( I + tX ) in a small neighborhood of t , thus motivating the parameterization of our statisticalmodel in (18).
Revised: ( I + tX ) in a small neighborhood of t , thus motivating the parameterization of our statistical model in (18).

Original: The next theorem gives a separation lower bound to ensure exact recovery of theclustering labels for Gaussian distributions.
Revised: The next theorem gives a separation lower bound to ensure exact recovery of the clustering labels for Gaussian distributions.

Original: Theorem 8 ( Exact recovery for clustering Gaussians) .
Revised: Theorem 8 ( Exact recovery for clustering Gaussians) .

Original: Let ∆ 2 :=
Revised: Let ∆ 2 :=

Original: min k ̸ = l d 2 ( V ( k ) , V ( l ) ) denotethe minimal pairwise separation among clusters, ¯ n := max k ∈
Revised: min k ̸ = l d 2 ( V ( k ) , V ( l ) ) denote the minimal pairwise separation among clusters, ¯ n := max k ∈

Original: [ K ] n k (and n := min k ∈ [ K ] n k ) themaximum (minimum) cluster size, and m :=
Revised: [ K ] n k (and n := min k ∈ [ K ] n k ) the maximum (minimum) cluster size, and m :=

Original: min k ̸ = l 2 n k n l n k + n l the minimal pairwise harmonic meanof cluster sizes.
Revised: min k ̸ = l 2 n k n l n k + n l the minimal pairwise harmonic mean of cluster sizes.

Original: Suppose the covariance matrix V i of Gaussian distribution ν i = N (0 , V i ) isindependently drawn from model (18) for i = 1 , 2 , . . . , n .
Revised: Suppose the covariance matrix V i of Gaussian distribution ν i = N (0 , V i ) is independently drawn from model (18) for i = 1 , 2 , . . . , n .

Original: Let β ∈ (0 , 1) .
Revised: Let β ∈ (0 , 1) .

Original: If the separation ∆ 2satisfies
Revised: If the separation ∆satisfies

Original: [Equation]
Revised: [Equation]

Original: 19) then the SDP (17) achieves exact recovery with probability at least 1 − C 2 n − 1 , provided that n ≥ C 3 log 2 n, t ≤ C 4 (cid:112) log n/ (cid:2) ( p + log ¯ n ) V 1 / 2 T 1 / 2 v (cid:3) , n/m ≤ C 5 log n, where V = max k (cid:13)(cid:13) V ( k ) (cid:13)(cid:13) op , T v = max 
Revised: then the SDP (17) achieves exact recovery with probability at least 1 − C 2 n − 1 , provided that [Equation] where V = max k (cid:13)(cid:13) V ( k ) (cid:13)(cid:13) op , T v = max k

Original: k Tr (cid:2)(cid:0) V ( k ) (cid:1) − 1 (cid:3) , and C i , i = 1 , 2 , 3 , 4 , 5 are constants.
Revised: Tr (cid:2)(cid:0) V ( k ) (cid:1) − 1 (cid:3) , and C i , i = 1 , 2 , 3 , 4 , 5 are constants.

Original: Remark 9 ( Further insight on pitfalls of barycenter-based Wasserstein K -means) .
Revised: Remark 9 ( Further insight on pitfalls of barycenter-based Wasserstein K -means) .

Original: Theorem 8suggests that different from Euclidean data, distributions after centering can be clustered if scales androtation angles vary (i.e., covariance-adjusted).
Revised: Theoremsuggests that different from Euclidean data, distributions after centering can be clustered if scales and rotation angles vary (i.e., covariance-adjusted).

Original: We further illustrate the rotation and scale effects onthe MNIST data that may mislead the centroid-based Wasserstein K -means, thus providing a realdata support for Example 3.
Revised: We further illustrate the rotation and scale effects on the MNIST data that may mislead the centroid-based Wasserstein K -means, thus providing a real data support for Example 3.

Original: Here we randomly sample two unbalanced clusters with 200 numbers of"0" and 50 numbers of "5".
Revised: Here we randomly sample two unbalanced clusters with 200 numbers of "0" and 100 numbers of "5".

Original: Fig.
Revised: Fig.

Original: 3 shows the clustering results for the centroid-based Wasserstein
Revised: 3 shows the clustering results for the centroid-based Wasserstein

Original: K -means and its oracle version where we replace the estimated barycenters µ 1 , µ 2 with the truebarycenters µ ∗ 1 , µ ∗ 2 computed on the true labels.
Revised: K -means and its oracle version where we replace the estimated barycenters µ 1 , µ 2 with the true barycenters µ ∗ 1 , µ ∗ 2 computed on the true labels.

Original: Comparing the Wasserstein distances W 2 ( µ 0 , µ ∗ 1 )and W 2 ( µ 0 , µ ∗ 2 ) , we see that the image µ 0 (containing digit "0") is closer to µ ∗ 2 (true barycenter ofdigit "5") and thus it cannot be classified correctly based on the nearest true barycenter (cf.
Revised: Comparing the Wasserstein distances W 2 ( µ 0 , µ ∗ 1 ) and W 2 ( µ 0 , µ ∗ 2 ) , we see that the image µ 0 (containing digit "0") is closer to µ ∗ 2 (true barycenter of digit "5") and thus it cannot be classified correctly based on the nearest true barycenter (cf.

Original: Fig.
Revised: Fig.

Original: 3 onthe left).
Revised: 3 on the left).

Original: Moreover, Wasserstein K -means based on estimated barycenters µ 1 , µ 2 yields two clustersof mixed "0" and "5".
Revised: Moreover, Wasserstein K -means based on estimated barycenters µ 1 , µ 2 yields two clusters of mixed "0" and "5".

Original: In both cases, the misclassification error is characterized by grouping similar degrees of angle and/or stretch.
Revised: In both cases, the misclassification error is characterized by grouping similar degrees of angle and/or stretch.

Original: Since there are two highly unbalanced clusters of distributions,
Revised: Since there are two highly unbalanced clusters of distributions,

Original: Wasserstein K -means is likely to enforce larger cluster to separate into two clusters and absorb thosearound centers (cf.
Revised: Wasserstein K -means is likely to enforce larger cluster to separate into two clusters and absorb those around centers (cf.

Original: Fig.
Revised: Fig.

Original: 3 on the right), leading to larger classification errors.
Revised: 3 on the right), leading to larger classification errors.

Original: We shall see thatin Section 4.3 the distance-based Wasserstein K -means and its SDP relaxation have much smallerclassification error rate on MNIST for the reason that we explained in Example 3 (cf.
Revised: We shall see that in Section 4.3 the distance-based Wasserstein K -means and its SDP relaxation have much smaller classification error rate on MNIST for the reason that we explained in Example 3 (cf.

Original: Lemma 4).
Revised: Lemma 4).

Original: ■
Revised: ■

Original:  Experiments 
Revised:  Experiments 4.1 Counter-example in Example 3 revisited

Original: 4.1 Counter-example in Example 3 revisited

Original: Our first experiment is to back up the claim about the failure of centroid-based Wasserstein Kmeans in Example 3 through simulations.
Revised: Our first experiment is to back up the claim about the failure of centroid-based Wasserstein Kmeans in Example 3 through simulations.

Original: Instead of using point mass measures that may resultsin instability for computing the barycenters, we use Gaussian distributions with small varianceas a smoothed version.
Revised: 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.

Original: We consider K = 2 , where cluster G ∗ 1 consists of m 1 many copies of( µ 1 , µ 2 ) pairs and m
Revised: We consider K = 2 , where cluster G ∗ 1 consists of m 1 many copies of ( µ 1 , µ 2 ) pairs and m

Original: 2 many µ 3 , and cluster G ∗ 2 consists of m 3 many copies of µ 4 .
Revised: 2 many µ 3 , and cluster G ∗ 2 consists of m 3 many copies of µ 4 .

Original: We chooseµ i as the following two-dimensional mixture of Gaussian distributions µ i = 0 .
Revised: We choose µ i as the following two-dimensional mixture of Gaussian distributions µ i = 0 .

Original: 5 N ( a i, 1 , Σ i, 1 ) +0 .
Revised: 5 N ( a i, 1 , Σ i, 1 ) + 0 .

Original: 5 N ( a i, 2 , Σ i, 2 ) for i = 1 , 2 , 3 , 4 .
Revised: 5 N ( a i, 2 , Σ i, 2 ) for i = 1 , 2 , 3 , 4 .

Original: Due to the space limit, detailed simulation setups and parametersare given in Appendix B. From Table 1, we can observe that Wasserstein SDP has achieved exactrecovery for all cases while barycenter-based Wasserstein K -means has only around 40% exactrecovery rate among all repetitions.
Revised: 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.

Original: In addition, Wasserstein SDP is more stable than distancebased Wasserstein K -means.
Revised: In addition, Wasserstein SDP is more stable than distancebased Wasserstein K -means.

Original: Denote ∆ k :
Revised: Denote ∆ k :

Original: = W 2 ( µ 3 , µ ∗ k ) as the squared distance between µ 3 andµ ∗ k for k = 1 , 2 , where µ ∗ k is the barycenter of G ∗ k .
Revised: = W 2 ( µ 3 , µ ∗ k ) as the squared distance between µ 3 and µ ∗ k for k = 1 , 2 , where µ ∗ k is the barycenter of G ∗ k .

Original: Let ∆ ∗ :
Revised: Let ∆ ∗ :

Original: = max k =1 , 2 max i,j ∈ G k W 2 ( µ i , µ j )and ∆ ∗ := min i ∈ G 1 ,j ∈ G 2 W 2 ( µ i , µ j ) be the maximum within-cluster distance and the minimumbetween-cluster distance respectively.
Revised: = max k =1 , 2 max i,j ∈ G k W 2 ( µ i , µ j ) and ∆ ∗ := min i ∈ G 1 ,j ∈ G 2 W 2 ( µ i , µ j ) be the maximum within-cluster distance and the minimum between-cluster distance respectively.

Original: From Table 9 in the Appendix, we can observe that ∆ ∗ < ∆ ∗ ,from which we can expect Wasserstein SDP to correctly cluster all data points in the Wasserteinspace.
Revised: From Table 7 in the Appendix, we can observe that ∆ ∗ < ∆ ∗ , from which we can expect Wasserstein SDP to correctly cluster all data points in the Wassertein space.

Original: Moreover, Table 1 shows that about 25% times that the distributions (as µ 3 ) in G ∗ 1 satisfy∆ 1 > ∆ 2 , implying those µ 3 to be likely assigned to the wrong cluster, which is consistent with
Revised: Moreover, Table 1 shows that about 25% times that the distributions (as µ 3 ) in G ∗ 1 satisfy ∆ 1 > ∆ 2 , implying those µ 3 to be likely assigned to the wrong cluster, which is consistent with

Original: Example 3.
Revised: Example 3.

Original: The experiment results also show that any copy of µ 3 is misclassified whenever exactrecovery fails for B-WKM, which means the misclassified rate for µ 3 equals to (1 − γ ) , where γ isthe exact recovery rate for B-WKM shown in Table 1.
Revised: The experiment results also show that any copy of µ 3 is misclassified whenever exact recovery fails for B-WKM, which means the misclassified rate for µ 3 equals to (1 − γ ) , where γ is the exact recovery rate for B-WKM shown in Table 1.

Original: Table 8 in the appendix further reports therun time comparison, from which we see that distance-based approaches are more computationallyefficient than the barycenter-based one in our settings.
Revised: Table 6 in the appendix further reports the run time comparison, from which we see that distance-based approaches are more computationally efficient than the barycenter-based one in our settings.

Original: [Table]
Revised: [Table]

Original: 4.2 Gaussian distributions
Revised: 4.2 Gaussian distributions

Original: Next, we simulate random Gaussian measures from model (18) with K = 4 and all cluster sizeequal.
Revised: Next, we simulate random Gaussian measures from model (18) with K = 4 and all cluster size equal.

Original: We set the centers of each cluster of Gaussians such that all pairwise distances among thebarycenters are all equal, i.e., W 22 ( N (0 , V ( k 1 ) ) , N (0 , V ( k 2 ) ))
Revised: We set the centers of each cluster of Gaussians such that all pairwise distances among the barycenters are all equal, i.e., W 22 ( N (0 , V ( k 1 ) ) , N (0 , V ( k 2 ) ))

Original: ≡ D for all k 1 , k 2 ∈ { 1 , 2 , 3 , 4 }with V = max k
Revised: ≡ D for all k 1 , k 2 ∈ { 1 , 2 , 3 , 4 } with V = max k

Original: ∥ V ( k ) ∥ op ∈ [4 . 5 , 5 . 5] .
Revised: ∥ V ( k ) ∥ op ∈ [4 . 5 , 5 . 5] .

Original: We fix the dimension p = 10 and vary the sample sizen = 200 , 400 , 600 .
Revised: We fix the dimension p = 10 and vary the sample size n = 200 , 400 , 600 .

Original: And we set the perturbation parameter t = 10 − 3 on the covariance matrix.
Revised: And we set the perturbation parameter t = 10 − 3 on the covariance matrix.

Original: Thesimulation results are reported over 100 times in each setting.
Revised: The simulation results are reported over 100 times in each setting.

Original: Fig.
Revised: Fig.

Original: 4 shows the misclassificationrate (log-scale) versus the squared distance D between centers.
Revised: 4 shows the misclassification rate (log-scale) versus the squared distance D between centers.

Original: We observe that when the distancebetween centers of clusters are larger than certain threshold (squared distance D > 10 − 3 in this case),then Wasserstein SDP can achieve exact recovery for different n, while the misclassification ratefor the two Wasserstein K -means are stably around 10% .
Revised: We observe that when the distance between centers of clusters are larger than certain threshold (squared distance D > 10 − 3 in this case), then Wasserstein SDP can achieve exact recovery for different n, while the misclassification rate for the two Wasserstein K -means are stably around 10% .

Original: But when the distance between centers ofclusters are relatively small, the two Wasserstein K -means  behave similarly or even slightly betterthan SDP.
Revised: And when the distance between centers of clusters are relatively small, the two Wasserstein K -means and SDP behave similarly.

Original: barycenter/distance-based Wasserstein K -means (B-WKM and D-WKM) for clustering Gaussians under n ∈ { 200 , 400 , 600 } .

Original: Due to the log-scale, 10 − 6 corresponds to exact recovery.

Original: 4.3 Real-data applications
Revised: 4.3 Real-data applications

Original: We consider three benchmark  datasets: MNIST, Fashion-MNIST and USPS handwriting digits.
Revised: We consider three benchmark image datasets: MNIST, Fashion-MNIST and USPS handwriting digits.

Original: Dueto complexity issues, we consider subsets of the whole datasets and randomly choose fixed number ofimages from each clusters based on 10 replicates for each cases.
Revised: 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.

Original: Here we used Sinkhorn divergence tocalculate the Wasserstein distance and the Bregman projection with 100 iterations for computing thebarycenters, which is efficient and stable for non-degenerate case in practice.
Revised: Here we used Sinkhorn divergence to calculate the Wasserstein distance and the Bregman projection with 100 iterations for computing the barycenters, which is efficient and stable for non-degenerate case in practice.

Original: For both Wasserstein K -means methods, we use the initialization method in analogue to the K -means++ for Euclideandata, i.e., the first cluster barycenter is chosen uniformly at random as one of the distributions,after which each subsequent cluster barycenter is chosen from the remaining distributions withprobability proportional to its squared Wasserstein distance from the distribution’s closest existingcluster barycenter. 
Revised: For both Wasserstein K -means methods, we use the initialization method in analogue to the K -means++ for Euclidean data, i.e., the first cluster barycenter is chosen uniformly at random as one of the distributions, after which each subsequent cluster barycenter is chosen from the remaining distributions with probability proportional to its squared Sinkhorn divergence from the distribution’s closest existing cluster barycenter. Codes using MATLAB and

Revised: Python implementing W-SDP and D-WKM are available at: https://github.com/Yubo02/

Revised: Wasserstein-K-means-for-clustering-probability-distributions .

Original: Further results and arguments as well as the new added discussion session areplaced at Appendix A due to page limit.
Revised: For MNIST dataset, we choose and fix two clusters: G ∗ 1 containing the number "0" and G ∗ 2 containing the number "5" for two cases. In Case 1, we randomly draw 200 number "0" and 100 number of "5" for each repetition. 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

Original: Checklist
Revised: K = 3 : G ∗ 1 containing the "T-shirt/top" (or the number "0" for USPS handwriting), G ∗ 2 containing the "Trouser" (or the number "5" for USPS handwriting) and G ∗ 3 containing the "Dress" (or the number "7" for USPS handwriting). The cluster sizes are unbalanced where we randomly choose 200 ,and 100 number from G ∗ 1 , G ∗ 2 and G ∗ 3 respectively. The error rates are shown in Table 2. Detailed setups and more results about F 1 scores (analogous to error rates) as well as time costs are placed at

Revised: Appendix A due to space limit.

Revised: From Table 2, we can observe that the performances for Wasserstein SDP (W-SDP) and distance-based

Revised: Wasserstein K -means (D-WKM) are better compared with barycenter-based Wasserstein K -means (B-WKM) for all cases.

Revised: And the results from Table 3 in Appendix A by using F 1 score are consistent with the results from Table 2.

Revised: The original K -means method behaves similar as barycenter-based

Revised: Wasserstein K -means in some cases and behaves less preferable for cases such as our experiment on

Revised: USPS handwriting digits.

Revised: In particular, the visualization of the clustering results for case 1 has been shown in Fig.

Revised: 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.

Revised: And this will lead to the high misclassification error for barycenter-based or centroid-based

Revised: Wasserstein K -means.

Revised: [Table]

Revised:  Discussion

Revised: In this paper

Revised: [Table]

Revised: we observed and analyzed the peculiar behaviors of Wasserstein barycenters and their results in clustering probability distributions.

Revised: After that

Revised: [Table]

Revised: we proposed the distance-based Kmeans approach (D-WKW) and its semidefinite program relaxation (W-SDP) by showing the exact recovery results for Gaussians theoretically and numerically.

Revised: For several real benchmark datasets, we showed results where D-WKM and W-SDP could outperform barycenter-based K-means approach (B-WKM).

Revised: 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.

Revised: 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.

Revised: One of our future goals would be addressing the corresponding computational complexity issues for real data.

Revised: Another goal is to find out more conclusive results where our approaches are preferable within or out of the realm of unbalanced clusters for real datasets.

Revised: Acknowledgments and Disclosure of Funding

Revised: Xiaohui Chen was partially supported by NSF CAREER grant DMS-1752614.

Revised: Yun Yang was partially supported by NSF grant DMS-2210717.

Original: The checklist follows the references.
Revised: The checklist follows the references.

Original: Please read the checklist guidelines carefully for information onhow to answer these questions.
Revised: Please read the checklist guidelines carefully for information on how to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
Revised: For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .

Original: You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
Revised: You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.

Original: For example:
Revised: For example:

Original: • Did you include the license to the code and datasets?
Revised: • Did you include the license to the code and datasets?

Original: [Yes] See Section ?? .•
Revised: [Yes] See Section ?? .

Original:  Did you include the license to the code and datasets?
Revised: • Did you include the license to the code and datasets?

Original: [No] The code and the data areproprietary.•
Revised: [No] The code and the data are proprietary.

Original:  Did you include the license to the code and datasets?
Revised: • Did you include the license to the code and datasets?

Original: [N/A]
Revised: [N/A]

Original: Please do not modify the questions and only use the provided macros for your answers.
Revised: Please do not modify the questions and only use the provided macros for your answers.

Original: Note that the
Revised: Note that the

Original: Checklist section does not count towards the page limit.
Revised: Checklist section does not count towards the page limit.

Original: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
Revised: In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes] (b) Did you describe the limitations of your work?

Original: [Yes] (c) Did you discuss any potential negative societal impacts of your work?[N/A]
Revised: [Yes] (c) Did you discuss any potential negative societal impacts of your work?

Original:  (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: [N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] (b) Did you include complete proofs of all theoretical results?
Revised: [Yes] (b) Did you include complete proofs of all theoretical results?

Original: [Yes] 
Revised: [Yes] 3.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [N/A] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: [N/A] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?[Yes]
Revised: [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original:  (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
Revised: [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: [N/A] 
Revised: [N/A] 4.

Original: 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?
Revised: 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?

Original: [Yes] (b) Did you mention the license of the assets?
Revised: [Yes] (b) Did you mention the license of the assets?

Original: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
Revised: [No] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

Original: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Original: [N/A] 
Revised: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: [N/A]
Revised: [N/A]