e-means：以十分之一 epoch 达到同等损失的更快 k-means

epoch 数量少一个数量级

在大量不同的规模和数据集上，e-means 在 1.2 到 3 个 epoch 内就达到了 Lloyd 在 25 个 epoch 时的测试损失。端到端实测，在 GPU 上相比最先进的 Lloyd 实现（Flash-KMeans^[2]）实现了 10–45 倍的挂钟时间缩减，在 CPU 上相比 Faiss^[1] 实现了 8–15 倍的缩减。在 k-means 中，速度和质量通常相互权衡，而 e-means 同时实现了更好的速度和损失。

e-means against Lloyd’s k-means, in epochs, CPU time and GPU time

k-means 问题

给定一个由 n 个点组成的数据集 X = {x₁, ..., x_n} ⊂ R^d 和一个整数 k，k-means 问题是要找到 k 个质心 C = (c₁, ..., c_k)，使每个点到其最近质心的平均平方距离最小：

\[ \mathrm{MSE}(X, C) = \frac{1}{n} \sum_{i=1}^{n} \|x_i - c_{a(i)}\|^2, \qquad a(i) = \arg\min_{j \in \{1,\dots,k\}} \|x_i - c_j\|^2 \]

\[ \mathrm{MSE}(X, C) = \frac{1}{n} \sum_{i=1}^{n} \|x_i - c_{a(i)}\|^2 \]

\[ a(i) = \arg\min_{j \in \{1,\dots,k\}} \|x_i - c_j\|^2 \]

这是一个静态数据集，不同于 VQ-VAE 内部的码本学习，后者的待聚类分布是不断变化的。

e-means 项目

e-means 的目标是为 k-means 问题提供一种现代优化规则。我们基于以下观察启动了该项目：

k-means 是机器学习中少数仍然广泛使用全批量优化的问题之一。
文献中存在小批量 k-means 方法^[3],[4],[5]，但大多基于 1967 年发布的 MacQueen 更新规则。
现代深度学习提供了可以增强小批量方法的优化工具。

对于 e-means，我们选择使用指数移动平均进行小批量更新^[6]，我们发现这对 k-means 问题最为高效。与以往的小批量方法不同，e-means 无需超参数调优，因为它通过缩放定律自动设置批量大小和步长（参见使用缩放定律的自动超参数调优）。此外，e-means 改进了初始化和死簇复活策略，增加了余弦调度的衰减率，并引入了一个收尾阶段，在训练结束时以可忽略的运行时开销降低训练损失。e-means 高度可并行化，并附带专用内核，将其算法优势转化为挂钟时间加速。

移动平均更新

每个小批量都会更新每个簇的计数和总和。质心是二者之比。

无需超参数调优

一切都由缩放定律推导而来。

开始与结束

新颖的初始化和收尾方法。运行结束时落在每个簇最后一个 epoch 的精确均值上。

余弦调度

步长 α 从 α₀ 退火到 0.1α₀，与学习率的做法相同。

死簇复活

大小低于阈值的簇通过拆分一个大簇来复活。

专用内核

GPU 内核融合了最近质心分配，将上述 epoch 节省转化为相应的挂钟时间加速。

质心更新规则

质心更新规则是 e-means 的核心组成部分。本节旨在概述 k-means 聚类的更新规则。所有方法都会为每个簇 j 维护分配给它的点的总和 S_j 和数量 N_j，并以二者之比作为质心：

\[ c_j = S_j / N_j \]

它们的区别仅在于 S_j 和 N_j 的更新方式。下文用 s_j 和 n_j 表示当前批次中分配给簇 j 的点的总和与数量（若没有点分配给它，则二者均为零）。

Lloyd 算法^[7] 将每个点分配给最近的质心，并以整个数据集作为批次从头重新计算统计量：

\[ S_j \leftarrow s_j, \qquad\quad N_j \leftarrow n_j \]

因此 c_j 是当前分配给它的点的精确均值。每次更新都需要对数据做一次完整遍历，该过程重复固定次数（通常为 20 到 50 次迭代）。

MacQueen 更新^[3] 会累积统计量而从不丢弃它们：

\[ S_j \leftarrow S_j + s_j, \qquad\quad N_j \leftarrow N_j + n_j \]

因此 c_j 是自训练开始以来分配给簇 j 的每个点的均匀均值，包括在较早的、可能较差的质心位置下分配的点。这一问题后来由 Staleness-Reduction Mini-Batch K-Means^[8] 部分解决，它是 k-means 与 MacQueen 之间的混合方法。它们获得了显著的加速，但仅在簇数量非常少时有效。

e-means 在小批次上运行，并用指数移动平均替代该求和：

\[ S_j \leftarrow (1-\alpha)\,S_j + \alpha\, s_j, \qquad\quad N_j \leftarrow (1-\alpha)\,N_j + \alpha\, n_j \]

\[ S_j \leftarrow (1-\alpha)\,S_j + \alpha\, s_j \]

\[ N_j \leftarrow (1-\alpha)\,N_j + \alpha\, n_j \]

因此较早的分配会以速率 α 被遗忘，质心跟踪的是近期的分配而非全部历史。该更新规则通常用于移动分布，因此被用于训练 VQ-VAE^[6]。我们发现，如果将其与 e-means 的其他组件结合，并适当调优超参数，它对静态数据集也极其高效。

为什么使用指数移动平均更新而不是梯度下降？

e-means 项目始于 Meta FAIR 的 Hugo Touvron 和 Hervé Jégou 的先前工作，他们发现通过梯度下降进行的小批次 k-means 可以胜过 Lloyd 算法。虽然我能够复现这些结果，但我发现基于梯度下降的质心更新在大学习率下性能不稳定，而指数移动平均更新则不然。我怀疑这是因为 EMA 更新保证质心始终位于分配给它的点的凸包内，而梯度下降更新可能会越过该凸包。

Test error against the effective step for exponential moving averages and plain gradient descent, at a large batch — 在大批量下，针对指数移动平均更新（EMA）和普通梯度下降（SGD），在每个尝试的步长设置下，测试误差与质心所采取的有效步长之间的对比。DEEP，100 万向量；SIFT，1000 万向量，16,384 个簇，5 个 epoch，3 个随机种子的中位数。纵轴为高于面板中最佳测试误差的百分比。在其最佳设置下，梯度下降最终比移动平均高出 1.6%（DEEP）和 0.4%（SIFT）。

为了在固定 epoch 数下达到最低损失，我们选择了指数移动平均更新，因此我们的方法命名为 e-means。

Lloyd 是 e-means 的一个特例：超参数研究

如果将批量设为整个数据集，B = n，并且 α = 1，那么 N_j 和 S_j 就恰好变成整个簇的计数和总和，因此一次 e-means 迭代就恰好变成一次 Lloyd 迭代。Lloyd 是 e-means 在没有学习率余弦调度下的极端情况。然而，我们的超参数搜索发现，这种极端设置高度次优，这解释了为什么 e-means 在很大程度上优于 Lloyd。下面，针对每个批量大小 B 和步长 α 测量损失：Lloyd 落在角落里，远在损失最低区域之外。

Validation loss for every batch size and constant step at 5 epochs; the lowest losses lie on a diagonal band, and Lloyd, at full batch and step 1, sits in the top-right corner 1.8% above the best — 每个批量 B 和步长 α 的验证损失，相对于最佳单元格，经过 5 个 epoch 后。DEEP，n = 10⁶，k = 1024。此处步长恒定（无余弦调度），因此 Lloyd 恰好是其中一个单元格：全批量，α = 1。最低损失沿对角线带分布，更大的批量需要更大的步长。

这项超参数研究促使我们推导出自动选择最优批量大小 B 和步长 α 的缩放定律。

e-means 更好地逃离局部极小值

为了更好地理解为什么 e-means 在收敛时比 Lloyd 达到更好的损失，我们在玩具数据集上绘制了几次训练，以评估它们逃离局部极小值的能力。在 200 个随机种子上重复训练，我们发现与 Lloyd 相比，e-means 陷入病态配置的可能性要小得多。

两个大小差异很大的团块，k = 8。Lloyd 将两个质心放在小团块上，而大团块则缺少质心。e-means 从相同的陷阱开始，并在最初的几个 epoch 内离开它，最终在大团块上有七个质心，小团块上有一个。在 200 次启动中，Lloyd 在 26% 的启动中保持卡住，而 e-means 在 3.5% 的情况下保持卡住。

3×3 网格中的九个团块，角部更重，k = 10。Lloyd 让中心团块没有质心，并将三个质心挤在一个角落里。e-means 陷入相同的排列，然后在第 10 个 epoch 左右爬出。在 200 次启动中，Lloyd 在 43% 的启动中保持卡住且差超过 1%，而 e-means 在 4% 的情况下差超过 1%。

围绕一个更重的中心团块的正方形上的四个团块，k = 5。Lloyd 将两个质心放在一个角落团块上，并将一个质心搁置在另外两个团块之间的间隙中，然后这两个团块共享它。e-means 从相同的起点开始，为每个团块分配一个质心。在 200 次启动中，Lloyd 在 15.5% 的启动中保持卡住，而 e-means 保持卡住的比例为 0%。

一条线上有五个数据团，k = 5。Lloyd 算法让中间一个轻量数据团没有质心，却给末端一个数据团分配了两个质心。e-means 从同一起点出发，给每个数据团各分配一个质心。在 200 次随机起点中，Lloyd 有 46% 的起点陷入停滞，而 e-means 的停滞比例为 0%。

我们的解释是，小批量提供了更多带噪声的更新，这些更新更有可能探索优化景观并逃离局部极小值。Lloyd 的损失始终下降，因此更倾向于利用已找到的局部极小值，而永远不会逃离它们。

利用缩放定律进行自动超参数调优

e-means 的一个关键特性是，与 Lloyd 算法相比，实践者无需设置任何额外的超参数。我们拟合了两条缩放定律，根据任意问题的特征 (n, k, T) 来设定批量大小 B 和步长 α₀，并发现它们在未见过的数据集和问题规模上表现出很强的泛化能力。

\[ B = c_B \, k^{b_k} \, n^{b_n} \, T^{b_T} \]

\[ \alpha_0 = \min\!\Big(c_\alpha \, k^{a_k} \, B^{a_B} \, n^{a_n} \, T^{a_T},\; 1\Big) \]

定律	常数	指数				遗憾值，留出数据集	遗憾值，留出规模	R²
定律	常数	k	B	n	T	遗憾值，留出数据集	遗憾值，留出规模	R²
步长 α₀	0.824	−0.46	0.90	−0.45	−0.11	0.07%	0.07%	0.95
批量 B，默认	0.035	0.13	不适用	0.88	0.81	0.19%	0.20%	0.96
批量 B，质量定律	0.0044	0.13	不适用	0.88	0.81	0.07%	0.06%	不适用

遗憾值 =（按定律得到的 MSE − 网格上的最佳 MSE）/ 最佳 MSE，取各配置的中位数。留出数据集：SIFT、SPACEV 以及 VIBE^[9] 数据集 arXiv 和 AG-News。留出规模：拟合时留出一个 n 值并进行预测。

我们提出两个方程来选择批量，即默认定律和质量定律。“质量定律”以吞吐量为代价，换取比“默认”定律更好的聚类质量。下面我们绘制缩放定律的训练数据：验证 MSE 随步长 α₀（左）和批量大小 B（右）的变化。

死簇复活的重要性

展示死簇复活重要性的玩具示例。

小批量 k-means 的主要问题是它容易产生“死”簇，即在优化过程中最终没有分配到任何点的质心。这会减小码本的有效规模，对于当今越来越多使用的大码本来说尤其成问题。

一种常见的缓解方法是重新初始化低于某个规模阈值的簇，如 EnCodec^[10] 中所做的那样，其中死质心被简单地替换为从传入批次中随机抽取的样本。虽然这有帮助，但并不能完全解决问题。我们不是从批次中采样任意点，而是通过分裂已有的大簇来复活死簇：我们按簇大小成比例的概率采样一个活跃簇，并将其分裂为两个。我们发现，这种改编自 Faiss k-means 的方法在消除死簇和保持码本容量方面明显更有效。

数据集	复活策略	死簇	码本占比	不平衡度	测试 MSE
DEEP	按比例分裂	47	0.3%	1.194	0.38762
	从批次中重采样（EnCodec）	2,343	14.3%	1.423	0.39124
	关闭	4,313	26.3%	1.647	0.39448
SIFT	按比例分裂	10	0.1%	1.186	0.19623
	从批次中重采样（EnCodec）	750	4.6%	1.233	0.19684
	关闭	2,080	12.7%	1.404	0.19797

100 万个向量，k = 16,384，5 个 epoch，对 5 个随机种子取平均；仅改变复活策略。如果最终分配给某个簇的向量少于十个（即平均簇大小的千分之一），则该簇被视为死簇。不平衡度：簇大小的不均匀程度（1 = 全部相等）。摘自论文的复活消融表。

大量实验验证

我们在 24 个不同的嵌入数据集和合成数据集及配置上验证了 e-means（见论文）。在视频中，我们针对 k-means 运行了 234 种 e-means 配置（16 个数据集，10 万到 1 亿个点）：在结果图中探索每条曲线。

Convergence curves vs Lloyd at 25 iterations — 训练（上）和测试（下）MSE，除以 25 个 epoch 后 Lloyd 的结果。从随机向量初始化。

Test error against wall-clock time on GPU, 10 million vectors, k=65,536 — 在单块 H100 上，1000 万个向量、65,536 个簇时测试 MSE 与墙钟时间的关系。普通 PyTorch Lloyd（数据在 GPU 上，bf16）的 25 个 epoch 需要 142 到 155 秒，而 Flash-KMeans^[2] 需要 14 到 22 秒。曲线上的标签标明了 epoch 数。

Epochs e-means needs to reach Lloyd's 25-epoch test error on the 11 VIBE embedding datasets: 1.6 to 2.4 — 在 VIBE 嵌入基准的 11 个数据集（文本、代码、图像和多模态嵌入，d = 200 到 3,072，k = 8,192 或 16,384）上，e-means 达到 Lloyd 在 25 个 epoch 中达到的测试误差所需的 epoch 数。预算间隔为 0.2 个 epoch，五个种子的中位数。每个数据集需要 1.6 到 2.4 个 epoch，而 Lloyd 需要 25 个。

库和代码

from faiss.contrib.e_means import Emeans # NumpyEmeans for NumPy em = Emeans(d, k, nepochs=T) # T epochs (5 by default) centroids, labels = em.train(x) # x: a PyTorch tensor (n, d) # labels: each vector's cluster in the last epoch

仓库：faiss/contrib/e_means。

引用本工作

如果您使用 e-means，请引用论文（arXiv:2611.09501）：

@misc{bessa2026emeansoutperformingkmeansin, title={e-means: outperforming k-means (in one-tenth the epochs)}, author={Swann Bessa and Pierre-Emmanuel Mazaré and Gergely Szilvasy and Hugo Touvron and David Picard and Matthijs Douze and Hervé Jégou}, year={2026}, eprint={2611.09501}, archivePrefix={arXiv}, primaryClass={cs.LG}, url={https://arxiv.org/abs/2611.09501}, }

参考文献

M. Douze, A. Guzhva, C. Deng, J. Johnson, G. Szilvasy, P.-E. Mazaré, M. Lomeli, L. Hosseini, H. Jégou. "The Faiss library." arXiv:2401.08281, 2024. ↑
S. Yang, H. Xi, Y. Zhao, M. Li, X. Fan, J. Zhang, H. Cai, Y. Lin, X. Li, K. Keutzer, S. Han, C. Xu, I. Stoica. "Flash-KMeans: Fast and Memory-Efficient Exact K-Means." arXiv:2603.09229, 2026. ↑
J. MacQueen. "Some Methods for Classification and Analysis of Multivariate Observations." Proceedings of the Fifth Berkeley Symposium on Mathematical Statistics and Probability, 1:281–297, 1967. ↑
D. Sculley. "Web-Scale K-Means Clustering." WWW, 2010. ↑
J. Newling, F. Fleuret. "Nested Mini-Batch K-Means." NeurIPS, 2016. ↑
A. van den Oord, O. Vinyals, K. Kavukcuoglu. "Neural Discrete Representation Learning." NeurIPS, 2017. ↑
S. Lloyd. "Least squares quantization in PCM." IEEE Transactions on Information Theory, 28(2):129–137, 1982 (circulated as a Bell Labs technical note since 1957). ↑
X. Zhu, J. Sun, Z. He, J. Jiang, Z. Wang. "Staleness-Reduction Mini-Batch K-Means." IEEE Transactions on Neural Networks and Learning Systems, 35(10):14424–14436, 2024. ↑
E. Jääsaari, V. Hyvönen, M. Ceccarello, T. Roos, M. Aumüller. "VIBE: Vector Index Benchmark for Embeddings." arXiv:2505.17810, 2025. ↑
A. Défossez, J. Copet, G. Synnaeve, Y. Adi. "High Fidelity Neural Audio Compression." Transactions on Machine Learning Research, 2023. arXiv:2210.13438. ↑