2. 核心思想
这篇论文的核心思想源于一个深刻的观察：在欧氏空间中等价的K-means的两种表述------基于质心（centroid-based）和基于距离（distance-based）------在Wasserstein空间中不再等价。
传统K-means聚类依赖于欧氏空间的两个关键性质：(1) 向量可以线性平均；(2) 空间是平坦的（零曲率）。这保证了著名的平行四边形法则（或其多点推广形式）成立，从而使得两种表述可以互换。
然而，Wasserstein空间 (P2(Rd),W2)(\mathcal{P}_2(\mathbb{R}^d), W_2)(P2(Rd),W2) 具有非负的Alexandrov曲率 ，是一个非平坦的测地空间。这导致Wasserstein重心（barycenter）表现出不规则性 （irregularity）和非鲁棒性（non-robustness）。因此，直接将欧氏K-means推广为基于Wasserstein重心的聚类方法（即用重心代替质心）会产生问题：重心可能无法很好地代表其所在簇内的数据点，甚至会导致错误的聚类分配。
基于此，论文提出并论证了：对于概率分布的聚类，应放弃基于重心的表述，转而采用更直接、更可靠的基于成对Wasserstein距离的聚类方法，并对其进行了半正定规划（SDP）松弛以获得理论保证和更好的性能。
3. 目标函数
论文定义了两种Wasserstein K-means的目标函数，并最终主推第二种。
(1) 基于重心的目标函数（论文认为其不可靠）
该目标函数模仿欧氏K-means的质心形式，但将欧氏距离替换为W2W_2W2距离，将样本均值替换为Wasserstein重心。
minν1,...,νK∈P2(Rd)∑i=1nmink∈[K]W22(μi,νk) \min_{\nu_1, ..., \nu_K \in \mathcal{P}2(\mathbb{R}^d)} \sum{i=1}^n \min_{k \in [K]} W_2^2(\mu_i, \nu_k) ν1,...,νK∈P2(Rd)mini=1∑nk∈[K]minW22(μi,νk)
其中 νk\nu_kνk 是第 kkk 个簇的重心。
(2) 基于距离的目标函数（论文主推的方法）
该目标函数直接推广欧氏K-means的成对距离形式，最小化每个簇内所有点对之间的平均 W22W_2^2W22 距离。
minG1,...,GK{∑k=1K1∣Gk∣∑i,j∈GkW22(μi,μj):⨆k=1KGk=[n]} \min_{G_1, ..., G_K} \left\{ \sum_{k=1}^K \frac{1}{|G_k|} \sum_{i,j \in G_k} W_2^2(\mu_i, \mu_j) \quad : \quad \bigsqcup_{k=1}^K G_k = [n] \right\} G1,...,GKmin⎩ ⎨ ⎧k=1∑K∣Gk∣1i,j∈Gk∑W22(μi,μj):k=1⨆KGk=[n]⎭ ⎬ ⎫
其中 GkG_kGk 是第 kkk 个簇的索引集合，⨆\bigsqcup⨆ 表示不相交并集。
4. 目标函数的优化过程
(1) 基于距离的Wasserstein K-means的贪心算法
论文提出了一种类似于Lloyd算法的迭代贪心算法来求解基于距离的目标函数：
- 初始化 ：给定初始聚类分配 Gk(1)G_k^{(1)}Gk(1)。
- 迭代更新 ：在第 ttt 次迭代中，对于每个分布 μi\mu_iμi，计算其到当前每个簇 Gk(t)G_k^{(t)}Gk(t) 的平均平方Wasserstein距离 ：
1∣Gk(t)∣∑s∈Gk(t)W22(μi,μs) \frac{1}{|G_k^{(t)}|} \sum_{s \in G_k^{(t)}} W_2^2(\mu_i, \mu_s) ∣Gk(t)∣1s∈Gk(t)∑W22(μi,μs) - 重新分配 ：将 μi\mu_iμi 分配给使其上述平均距离最小的簇：
Gk(t+1)={i∈[n]:1∣Gk(t)∣∑s∈Gk(t)W22(μi,μs)≤1∣Gj(t)∣∑s∈Gj(t)W22(μi,μs), ∀j∈[K]} G_k^{(t+1)} = \left\{ i \in [n] : \frac{1}{|G_k^{(t)}|} \sum_{s \in G_k^{(t)}} W_2^2(\mu_i, \mu_s) \leq \frac{1}{|G_j^{(t)}|} \sum_{s \in G_j^{(t)}} W_2^2(\mu_i, \mu_s), \, \forall j \in [K] \right\} Gk(t+1)=⎩ ⎨ ⎧i∈[n]:∣Gk(t)∣1s∈Gk(t)∑W22(μi,μs)≤∣Gj(t)∣1s∈Gj(t)∑W22(μi,μs),∀j∈[K]⎭ ⎬ ⎫ - 重复步骤2-3直到收敛。
(2) SDP松弛的优化过程
由于直接求解组合优化问题是NP-hard的，论文采用了SDP松弛。
- 矩阵重构 ：定义距离矩阵 A∈Rn×nA \in \mathbb{R}^{n \times n}A∈Rn×n，其中 Aij=W22(μi,μj)A_{ij} = W_2^2(\mu_i, \mu_j)Aij=W22(μi,μj)。将聚类分配表示为一个二元指示矩阵 H∈{0,1}n×KH \in \{0,1\}^{n \times K}H∈{0,1}n×K，其中 Hik=1H_{ik}=1Hik=1 当且仅当 i∈Gki \in G_ki∈Gk。
- 整数规划 ：基于距离的目标函数可以等价地写成一个非线性0-1整数规划：
min{⟨A,HBH⊤⟩:H∈{0,1}n×K,H1K=1n} \min \left\{ \langle A, H B H^\top \rangle : H \in \{0, 1\}^{n \times K}, H\mathbf{1}_K = \mathbf{1}_n \right\} min{⟨A,HBH⊤⟩:H∈{0,1}n×K,H1K=1n}
其中 B=diag(∣G1∣−1,...,∣GK∣−1)B = \text{diag}(|G_1|^{-1}, ..., |G_K|^{-1})B=diag(∣G1∣−1,...,∣GK∣−1)。 - 变量替换与松弛 ：令 Z=HBH⊤Z = H B H^\topZ=HBH⊤。ZZZ 是一个对称半正定矩阵，满足 Tr(Z)=K\text{Tr}(Z) = KTr(Z)=K, Z1n=1nZ\mathbf{1}_n = \mathbf{1}nZ1n=1n, 且 Z≥0Z \geq 0Z≥0（元素非负）。将这些凸约束保留，而舍弃 ZZZ 必须由 HHH 生成的非凸约束，得到SDP松弛问题：
minZ∈Rn×n{⟨A,Z⟩:Z⊤=Z,Z⪰0,Tr(Z)=K,Z1n=1n,Z≥0} \min{Z \in \mathbb{R}^{n \times n}} \left\{ \langle A, Z \rangle : Z^\top = Z, Z \succeq 0, \text{Tr}(Z) = K, Z\mathbf{1}_n = \mathbf{1}_n, Z \geq 0 \right\} Z∈Rn×nmin{⟨A,Z⟩:Z⊤=Z,Z⪰0,Tr(Z)=K,Z1n=1n,Z≥0} - 求解 ：该SDP是一个凸优化问题，可以使用标准的凸优化求解器（如内点法）进行求解。求解得到的 Z∗Z^*Z∗ 通常是一个实值矩阵，需要通过后续的舍入（rounding）步骤（例如K-means对 Z∗Z^*Z∗ 的行进行聚类）来恢复最终的离散聚类标签。
5. 主要贡献点
- 揭示了Wasserstein重心的缺陷：通过理论分析（平行四边形法则失效）和具体反例（如论文中的Example 1, 2, 3），系统性地论证了基于重心的Wasserstein K-means在不规则性和非鲁棒性方面存在根本性问题，可能导致不合理的、对扰动敏感的聚类结果。
- 推广了基于距离的Wasserstein K-means：明确提出了适用于Wasserstein空间的、更稳健的基于成对距离的聚类目标函数，并设计了相应的贪心算法。
- 建立了理论保证 ：针对高斯分布聚类这一重要特例，证明了所提出的SDP松弛方法在簇间分离度足够大时，能够以高概率实现精确恢复（exact recovery）的真实聚类标签。
- 全面的实证验证：通过精心设计的模拟实验（验证反例）和在MNIST、Fashion-MNIST、USPS等真实图像数据集上的实验，一致地证明了基于距离的方法（及其SDP松弛）在分类性能上显著优于基于重心的方法。
6. 算法的实现过程详解
论文中涉及三种主要算法的实现：
**(1) 基于重心的Wasserstein K-means **(B-WKM)
- 初始化：采用类似K-means++的策略，随机选择一个分布作为第一个重心，后续重心根据其到现有重心的Sinkhorn散度（Wasserstein距离的正则化近似）进行概率采样。
- 迭代 ：
- 分配步骤 ：计算每个分布 μi\mu_iμi 到所有重心 νk\nu_kνk 的 W2W_2W2 距离（或Sinkhorn散度），将其分配给最近的重心。
- 更新步骤 ：对每个簇 GkG_kGk，求解其Wasserstein重心 νk=argminν1∣Gk∣∑i∈GkW22(μi,ν)\nu_k = \arg\min_\nu \frac{1}{|G_k|} \sum_{i \in G_k} W_2^2(\mu_i, \nu)νk=argminν∣Gk∣1∑i∈GkW22(μi,ν)。这本身是一个复杂的优化问题，论文中使用了迭代Bregman投影（IBP）算法，并限制了100次迭代。
- 输出：最终的聚类分配和重心。
**(2) 基于距离的Wasserstein K-means **(D-WKM)
- 初始化：同样采用K-means++的变体，但基于成对距离进行初始化。
- 迭代 ：
- 分配步骤 ：对于每个 μi\mu_iμi，计算其到每个簇 GkG_kGk 内所有成员的平均 W22W_2^2W22 距离，并分配给平均距离最小的簇。无需显式计算重心。
- 更新步骤：仅更新簇成员集合。
- 输出：最终的聚类分配。
**(3) Wasserstein SDP **(W-SDP)
- 预计算 ：首先计算完整的 n×nn \times nn×n 成对 W22W_2^2W22 距离矩阵 AAA（使用Sinkhorn散度近似）。
- 构建SDP ：根据矩阵 AAA 和聚类数 KKK，构建前述的SDP松弛问题。
- 求解SDP ：使用凸优化工具箱（论文代码使用了MATLAB或Python）求解该SDP问题，得到最优解 Z∗Z^*Z∗。
- 舍入 ：对矩阵 Z∗Z^*Z∗ 的每一行（代表一个数据点）进行标准的K-means聚类，从而将实值解 Z∗Z^*Z∗ 转化为离散的聚类标签。
- 输出：舍入后的聚类标签。
7. 局限性分析与总结
尽管这篇论文做出了重要贡献，但其方法也存在一些明显的局限性，论文作者在第5节"Discussion"中也坦诚地指出了部分问题：
-
计算复杂度高，可扩展性差：
- 距离矩阵计算 ：无论是D-WKM还是W-SDP，都需要计算并存储一个 O(n2)O(n^2)O(n2) 的成对距离矩阵。对于大规模数据集（nnn 很大），这在时间和空间上都是不可行的。
- SDP求解 ：求解一个 n×nn \times nn×n 的SDP问题的计算复杂度通常为 O(n6)O(n^6)O(n6) 或更高，这严重限制了W-SDP方法在大规模场景下的应用。作者也明确将解决计算复杂度问题列为未来工作目标。
-
Wasserstein距离计算的近似：
- 论文中使用了Sinkhorn散度 来近似计算 W2W_2W2 距离。虽然Sinkhorn算法计算效率高且数值稳定，但它引入了正则化偏差，得到的并非精确的Wasserstein距离。这种近似可能会影响聚类结果，尤其是在对距离精度要求很高的场景。
-
理论保证的局限性：
- 论文中关于SDP精确恢复的理论保证（Theorem 8）仅针对均值为零的高斯分布（即只考虑协方差矩阵的差异）。对于更一般的概率分布（如非高斯、有非零均值、高维稀疏支撑等），尚缺乏类似的理论保证。
-
SDP舍入的不确定性：
- SDP松弛提供了一个凸的下界，但最终的聚类结果依赖于舍入步骤（如对 Z∗Z^*Z∗ 做K-means）。这个舍入过程本身是启发式的，没有理论保证能完美恢复整数解，尤其是在信噪比较低、簇间分离度不够大的情况下。
-
对不平衡簇的敏感性：
- 尽管论文指出基于距离的方法在不平衡簇（如MNIST实验中"0"和"5"的数量不等）上表现更好，但其目标函数 1∣Gk∣∑i,j∈GkW22(μi,μj)\frac{1}{|G_k|} \sum_{i,j \in G_k} W_2^2(\mu_i, \mu_j)∣Gk∣1∑i,j∈GkW22(μi,μj) 本身就对簇大小敏感。非常小的簇可能会因为其内部距离的方差较大而难以被正确识别。
总结：这篇论文是一篇高质量的工作，它通过深刻的几何洞察指出了将传统聚类方法盲目推广到Wasserstein空间的陷阱，并提出了一个更稳健、性能更好的替代方案。其核心贡献在于概念上的澄清和理论验证。然而，其提出的最优方法（W-SDP）受限于高昂的计算成本，这使得它目前更适合作为一个理论基准或用于中小规模问题。未来的研究方向很可能是开发能够兼顾性能和可扩展性的近似算法，例如利用随机投影、子采样（如作者另一工作Sketch-and-lift）或开发针对Wasserstein距离的高效谱聚类方法。