
性能突破:速度与质量兼得
在广泛的规模和数据集测试中,e-means 仅需 1.2 到 3 个训练轮次(epochs),其测试损失即可低于 Lloyd 算法在 25 轮时的水平。端到端测量显示,与最先进的 Lloyd 实现(Flash-KMeans)相比,e-means 在 GPU 上的实际运行时间缩短了 10–45 倍;与 Faiss 相比,CPU 上的运行时间缩短了 8–15 倍。传统 k-means 算法往往需要在速度和质量之间进行权衡,而 e-means 成功打破了这一局限,同时实现了更快的速度和更低的损失。
核心机制:指数移动平均与自动调参
k-means 问题的目标是在静态数据集中找到 k 个质心,以最小化每个点到其最近质心的均方距离。e-means 旨在为这一问题提供现代化的优化规则,其核心观察在于:k-means 是机器学习中少数仍广泛使用全批量(full-batch)优化的问题之一,而现代深度学习中的小批量优化工具尚未被充分应用。
e-means 采用指数移动平均(EMA)进行小批量更新,取代了传统的求和更新规则。在这种机制下,旧的分配以速率 α 被遗忘,质心跟踪近期的分配而非整个历史。这种方法不仅具有高度并行性,还通过专用内核将算法优势转化为实际的运行加速。
与以往需要复杂超参数调整的小批量方法不同,e-means 通过拟合两个“缩放定律”,根据问题特征(数据点数量 n、簇数量 k、训练轮次 T)自动设置批量大小 B 和初始步长 α₀。这种自动化机制在未见过的数据集和问题规模上表现出强大的泛化能力,从业者无需手动干预。
优化策略:逃离局部极小与死簇复活
实验表明,e-means 比 Lloyd 算法更能有效逃离局部极小值。小批量更新引入的噪声有助于探索优化景观,而 Lloyd 算法由于损失总是递减,更容易陷入病态配置且无法逃脱。
针对小批量 k-means 常见的“死簇”问题(即某些质心在优化过程中未分配到任何点),e-means 改进了复活策略。不同于从输入批次中随机采样替换死质心,e-means 通过分裂大型现有簇来复活死簇:以与大小成正比的概率采样一个活跃簇并将其一分为二。这种方法在消除死簇和保留码本容量方面显著更有效。
项目背景与应用
e-means 项目始于 Meta FAIR 此前由 Hugo Touvron 和 Hervé Jégou 开展的工作。研究发现,基于梯度下降的小批量 k-means 虽可优于 Lloyd 算法,但在大学习率下表现不稳定;而 EMA 更新能保证质心始终位于分配给它的点的凸包内,因此更为稳定高效。
目前,e-means 已集成至 Faiss 库中,开发者可通过 faiss.contrib.e_means 直接调用。该算法已在 24 种不同的嵌入和合成数据集及配置上得到验证,并在包含 10 万至 1 亿个点的 16 个数据集上进行了大规模配置测试。