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

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress
Revised: 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 objectsbased 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 relieson 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: Inthe 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 ariseas probability distributions and a natural generalization to handle measure-valueddata 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 Alexandrovcurvature 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 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.

Original: In the special caseof clustering Gaussian distributions, we show that the SDP relaxed Wasserstein
Revised: In the special caseof 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 underthe 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 demonstratethat distance-based K -means can achieve better classification performance overthe standard centroid-based K -means for clustering probability distributions andimages.

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 analysisin 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 metricspace ( 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 recoverthe 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 manyapplications [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 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.

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: Normalizingthe 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 wouldlose 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 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,

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

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 isan iterative procedure alternating the partition and centroid estimation steps.

Original: Specifically, given aninitial centroid estimate β (1)1 , . . .
Revised: Specifically, given aninitial 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 thet -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) alternates 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 thefollowing 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,
Revised: 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,

Original: [Equation]
Revised: [Equation]

Original: and X i ∈ R
Revised: and X i ∈ R

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 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 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 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 ) isdefined as

Original: [Equation]
Revised: [Equation]

Original: (6) where minimization over γ runs over all possible couplings with marginals µ and ν .
Revised: (6) 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
Revised: It is well-knownthat the Wasserstein space is a metric space (in fact a geodesic space) with non-negative curvature inthe Alexandrov sense

Original: [Lott, 2008].
Revised: [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 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.

Original: Existing work.
Revised: Existing work.

Original: Since the K -means clustering is a worst-case NP-hard problem [Aloise et al.
Revised: Since the K -means clustering is a worst-case NP-hard problem [Aloise et al.

Original: ,2009], approximation algorithms have been extensively studied in literature including: Llyod’salgorithm
Revised: ,2009], approximation algorithms have been extensively studied in literature including: Llyod’salgorithm

Original: [Lloyd, 1982], spectral methods [von Luxburg, 2007, Meila and Shi, 2001, Ng et al.
Revised: [Lloyd, 1982], spectral methods [von Luxburg, 2007, Meila and Shi, 2001, Ng et al.

Original: ,2001], semidefinite programming (SDP) relaxations
Revised: ,2001], semidefinite programming (SDP) relaxations

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

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

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

Original: More discussions of
Revised: More discussions of

Original: Wasserstein K -means using barycenters can be found in Section 2.1.
Revised: Wasserstein K -means using barycenters can be found in Section 2.1.

Original: Verdinelli and Wasserman[2019] proposed a modified Wasserstein distance for distribution clustering.
Revised: Verdinelli and Wasserman[2019] proposed a modified Wasserstein distance for distribution clustering.

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 isto 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 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.

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, 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.

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 Wassersteinspace has been recently considered by Domazakis et al.

Original: Specifically, it is an iterativealgorithm proceeds as following.
Revised: Specifically, it is an iterativealgorithm 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 assignseach probability measure µ 1 , . . . , µ n to its nearest centroid in the Wasserstein geometry at the t -thiteration 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 Wassersteinanalog 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-changeablyuse 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: 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 .

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, 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.

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, theprobability 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 L 0and 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 ) / 2and µ 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 optimaltransport 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 att = 0 .

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

Original: 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 showsthat 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 turnout to be essential in centroid-based Wasserstein K -means and irregularity of the barycenter may failto represent the cluster (see more details in Example 3 and Remark 9 below), this counter-example isour motivation to seek alternative formulation.

Original: [Figure]
Revised: [Figure]

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 contributingprobability 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 onthe 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 seethat the optimal transport map T :

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

Original: [Equation]
Revised: [Equation]

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

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

Original: This non-robustness can be attribute to the discontinuity
Revised: This non-robustness can be attribute to the discontinuity

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

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 tounreasonable 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 maysignificantly slow down the convergence of commonly used barycenter approximating algorithmssuch 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 ofsuch 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 thisexample 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, 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.

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 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) .

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

Original: R

Original: [Equation]
Revised: [Equation]

Original: there is always some point
Revised: 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 Xthan 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 )
Revised: 2 highlighted by a red circle,consisting of m copies of ( µ 1 , µ 2 ) pairs and one copy of µ 3 )

Original: ; the second cluster G ∗ 2 containingcopies of µ 4 is highlighted by a blue circle.
Revised: ; 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 twopoints, where the two supporting points are connected by a dashed line for easy illustration.

Original: Morespecifically, we set
Revised: Morespecifically, 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 ) satisfiesy 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 [Equation] [Table] where µ ∗ k denotes the Wasserstein barycenter of cluster G k for k = 1 , 2 .
Revised: 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 µ 2is ˜ µ ∗ 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 µ ∗ 1of 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 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 .

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 µ 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.

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

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 smallchange in the orientation of distribution µ 1 may completely alter the clustering membership of µ 3based 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 , thenthe 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 verticalaxis.

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 tothe 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 canbe 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 alwaysassigns µ 3 into cluster G ∗ 2 no matter x

Original: <
Revised: <

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

Original: [Figure]
Revised: ■

Original: [Equation] [Figure] 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 insteadpropose 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’salgorithm 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 membershipestimate 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 theaveraged squared W 2 distances to all current members in every cluster, leading to an updated clustermembership 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 thatthe center-based and distance-based Wasserstein K -means formulations may not necessarily beequivalent 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 exampleillustrating 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 ,
Revised: If µ i = N ( m i , V i ) with positive-definite covariance matricesΣ i ≻ 0 ,

Original: then the squared 2 -Wasserstein distance can be expressed as the sum of the squared Euclideandistance on the mean vector and
Revised: then the squared 2 -Wasserstein distance can be expressed as the sum of the squared Euclideandistance on the mean vector and

Original: [Equation]
Revised: [Equation]

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 todenote the unique symmetric square root matrix of V ≻ 0 .

Original: That is,
Revised: That is,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: 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 withcumulative distribution function (cdf) F i , then the Wasserstein distance can be written in terms of thequantile transform

Original: [Equation]
Revised: [Equation]

Original: [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 Wassersteinspace 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 largeor 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 solvingthe 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 solvingthe 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 isto 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 SDPrelaxation 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
Revised: , X n ∈ R p are generatedfrom a Gaussian mixture model with isotropic known variance

Original: [Chen and Yang, 2021].
Revised: [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 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.

Original: Therefore,we will derive a Wasserstein SDP relaxation to the combinatorial optimization problem (4) using the squared distance matrix A n × n = { a ij } with a ij = W 22 ( µ i , µ j ) .
Revised: Therefore,we will derive a Wasserstein SDP relaxation to the combinatorial optimization problem (4) usingthe squared distance matrix A n × n = { a ij } with a ij = W 22 ( µ i , µ j ) .

Original: Concretely, we can one-to-onereparameterize any partition ( G 1 , . . .
Revised: Concretely, we can one-to-onereparameterize 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 0-1integer program,

Original: [Equation]
Revised: [Equation]

Original: where 1 n is the n × 1 vector of all ones and B = diag ( | G 1 | − 1 , . .
Revised: 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 variableto 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 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.

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 distributionssince optimal separation conditions for exact recovery based on the Euclidean mean componenthave 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, andthe 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 randommatrix 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 ) ispsd 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 geodesicemanating 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 statisticalmodel 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 theclustering 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 ) ) denotethe 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 ) themaximum (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 meanof 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 ) isindependently 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 ∆ 2satisfies

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 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: (19) then the SDP (17) achieves exact recovery with probability at least 1 − C 2 n − 1 , provided that [Equation] 5 log n, where V = max k (cid:13)(cid:13) V ( k ) (cid:13)(cid:13) op , T v = max k 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: Theorem 8suggests that different from Euclidean data, distributions after centering can be clustered if scales androtation 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 onthe MNIST data that may mislead the centroid-based Wasserstein K -means, thus providing a realdata 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 50 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 truebarycenters µ ∗ 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 ofdigit "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 onthe 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 clustersof 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 similardegrees 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 thosearound 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 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.

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

Original: ■
Revised: ■

Original:  Experiments
Revised:  Experiments4.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 resultsin instability for computing the barycenters, we use Gaussian distributions with small varianceas 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 parametersare given in Appendix A. 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.

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 minimumbetween-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 5 in the Appendix, we can observe that ∆ ∗ < ∆ ∗ ,from which we can expect Wasserstein SDP to correctly cluster all data points in the Wasserteinspace.

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

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 4 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.

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 sizeequal.

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 thebarycenters 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 sizen = 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: Thesimulation 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 misclassificationrate (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 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% .

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

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 A real-data application

Original:  We consider three benchmark datasets: MNIST, Fashion-MNIST and USPS handwriting digits.
Revised: Finally, we run our Wasserstein SDP algorithm against Wasserstein K -means on the MNIST dataset. We choose two clusters: G ∗ 1 containing the number "0" and G ∗ 2 containing the number "5", sothat the number of clusters is K = 2 in the algorithms. The cluster sizes are unbalanced with|

Original: Dueto complexity issues, weconsider subsets of the whole datasets and randomly choose fixed number ofimages from each clusters based on 10 replicates for each cases.
Revised: G ∗ 1 | / | G ∗ 2 | = 4 , where we randomly choose 200 number "0" and 50 number of "5" for each repetition. The results are shown in Table 2 based on 10 replicates.

Original: Here we usedSinkhorn 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 the Bregman projectionwith 100 iterations for computing the barycenters, which is efficient and stable for non-degeneratecase 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 inanalogue to the K -means++ for Euclidean data, i.e., the first cluster barycenter is chosen uniformlyat random as one of the distributions, after which each subsequent cluster barycenter is chosen fromthe remaining distributions with probability proportional to its squared Wasserstein distance from thedistribution’s closest existing cluster barycenter.

Revised: From Table 2 we can see that the performances for

Revised: Wasserstein SDP (W-SDP) and distance-based Wasserstein K -means (D-WKM) are better comparedwith barycenter-based Wasserstein K -means (B-WKM).

Original: Further results and arguments as well as the new added discussion session areplaced at Appendix A due to page limit.
Revised: [Table] The visualization of the clustering results has been shown in Fig. From this figure we can find thatthe classification criterion for B-WKM will end up with the closeness to certain shape of "0", which ischaracterized by certain angle or the degree of stretch. And this will lead to the high misclassificationerror for barycenter-based or centroid-based Wasserstein K -means.

Original: Checklist
Revised: Checklist

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 onhow 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 referencingthe 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 areproprietary.•

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 instructionsblock and only keep the Checklist section heading above along with the questions/answers below.

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes](c)

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: Did you discuss any potential negative societal impacts of your work?

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

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

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

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

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

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?
Revised: mental results (either in the supplemental material or as a URL)?

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

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

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., type
Revised: (d) Did you include the total amount of compute and the type of resources used (e g., type

Original: of GPUs, internal cluster, or cloud provider)?
Revised: of GPUs, internal cluster, or cloud provider)?

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

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]
Revised: [Yes]

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

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

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

Original: [No]
Revised: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

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

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: information or offensive content?
Revised: information or offensive content?

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

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
Revised: (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?

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

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

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

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?
Revised: spent on participant compensation?

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