[Paper] BalLOT：平衡 k-means 聚类与最优传输

Source: arXiv - 2512.05926v1

Overview

BalLOT（Balanced k‑means via Optimal Transport）解决了聚类中的一个经典难点：在最小化簇内方差的同时，强制每个簇包含大致相同数量的点。通过将分配步骤表述为最优传输（optimal‑transport）问题，作者设计了一种交替最小化算法，既快速又在具有挑战性的合成数据集上具有可证明的可靠性。

Key Contributions

• 新算法（BalLOT）：将最优传输耦合引入 k‑means 的交替最小化循环，以实现簇大小的平衡。
• 理论保证：
  • 证明对一般数据而言，每一次迭代都会产生 整数（即硬分配）传输计划，消除了后处理的需求。
  • 提供了景观分析，展示在随机球模型下能够精确或部分恢复植入的簇。
• 初始化策略：当数据满足轻度分离条件时，可实现对真实簇的 一步 恢复。
• 大量实证验证：展示了相较于现有平衡 k‑means 基线在速度和聚类质量上的优势。

Methodology

BalLOT 采用熟悉的两步 k‑means 流程——分配和中心更新——但用 最优传输（OT）问题 替代了标准的最近中心分配：

1. 通过 OT 进行分配

   • 将当前中心视为一组“供给”点，每个点的容量等于期望的簇大小（例如，完美平衡簇的容量为 (n/k)）。
   • 将数据点视为“需求”节点。
   • 求解 线性指派问题（OT 的特例），在满足容量约束的前提下最小化总平方欧氏距离。该问题的解得到一个 耦合 矩阵，直接编码了每个点所属的簇。

2. 中心更新

   • 根据 OT 步得到的硬分配，像经典 k‑means 那样将每个中心重新计算为其分配点的均值。

3. 交替循环

   • 重复 OT 分配和中心更新，直至收敛。

由于 OT 子问题是 平衡指派（经典的 Hungarian‑type 问题），其最坏情况时间复杂度为 (O(n^3))，但作者利用稀疏性和热启动实现了接近线性的实际性能。

Results & Findings

• 速度：在规模达 10 万点、50 个簇的合成数据集上，BalLOT 在 10 秒以内收敛，快于最先进的平衡 k‑means 求解器（后者常需数分钟）。
• 聚类质量：通过簇内平方和（WCSS）和簇平衡误差衡量，BalLOT 始终取得更低的 WCSS 且完美满足平衡约束，区别于那些为平衡而牺牲目标函数质量的启发式后处理方法。
• 恢复保证：在随机球模型（点来自分离良好的各向同性高斯球）下，作者证明 BalLOT 能以高概率恢复真实划分，即使分离度仅为 (O(\sqrt{\log k}))。
• 一步恢复：使用文中提出的初始化（例如，以少量相距遥远的点作为种子），BalLOT 能在第一次 OT 步后即正确分配所有点，省去多次迭代的需求。

Practical Implications

• 公平资源分配：在工作负载必须均匀分配的系统中（如数据库分片、微服务负载均衡），BalLOT 提供了一种在满足容量限制的前提下进行数据划分的原则性方法。
• 平衡数据抽样：在需要平衡小批量训练机器学习模型的场景（如联邦学习或类别不平衡分类），BalLOT 可生成平衡簇作为批次种子。
• 可扩展的预处理：由于 OT 步可以借助现代 GPU 求解器或近似方法（如 Sinkhorn 缩放）加速，BalLOT 适用于大规模流水线，传统平衡 k‑means 往往成为瓶颈。
• 可解释性：整数耦合保证每次迭代都有硬分配，使得得到的簇可直接使用，无需额外的取整或后处理。

Limitations & Future Work

• 可扩展性上限：尽管实际运行时间表现良好，精确的 OT 步仍呈超线性增长；探索近似 OT（如熵正则化）有望将 BalLOT 推向百万级点。
• 等簇大小假设：当前形式强制严格相等；将其扩展到 软 平衡约束（允许一定容差）将提升适用范围。
• 对异常值的鲁棒性：理论分析基于干净的随机球数据；处理重尾噪声或对抗性异常值仍是未解挑战。
• 真实世界基准：本文侧重合成实验；在特定领域数据集上评估 BalLOT（如图像块聚类、网络流量划分）是自然的后续工作。

总体而言，BalLOT 将最优传输理论与实用聚类相结合，为开发者提供了一种快速、平衡且具备坚实理论支撑的 k‑means 替代方案。

Authors

• Wenyan Luo
• Dustin G. Mixon

Paper Information

• arXiv ID: 2512.05926v1
• Categories: stat.ML, cs.DS, cs.IT, cs.LG, math.OC
• Published: December 5, 2025
• PDF: Download PDF