E-means：同等损失下 k-means 速度快 10 倍

少一个数量级的历元 在大范围的尺度和数据集上，e-means 在 25 个历元（1.2 到 3 个历元）中达到了低于 Lloyd's 的测试损失。端到端测量发现，与最先进的 Lloyd 实现（Flash-KMeans [2]）相比，GPU 上的挂钟减少了 10-45 倍，与 Faiss [1] 相比，CPU 上的挂钟减少了 8-15 倍。 k-means 中的速度和质量通常是相互矛盾的，而 e-means 可以实现更好的速度和损失。 k 均值问题 给定 n 个点的固定数据集 X = {x 1 , ..., x n } ⊂ R d 和整数 k ，k 均值问题是找到 k 个质心 C = (c 1 , ..., 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 均值方法 [3] 、 [4] 、 [5] ，但大多基于 1967 年发布的 MacQueen 更新规则。现代深度学习提供了可以增强小批量方法的优化工具。对于 e 均值，我们选择使用指数移动平均值 [6] 进行小批量更新，我们发现这对于 k 均值问题最有效。与以前的小批量方法相反，e-means 不需要超参数调整，因为它使用缩放法则自动设置批量大小和步长（请参阅使用缩放法则自动调整超参数）。此外，e-means 改进了初始化和死簇复兴策略，添加了余弦调度衰减率，并引入了最终确定阶段，该阶段可减少训练损失，在训练结束时运行时间损失可以忽略不计。 e-means 具有高度可并行性，并配有专用内核，可将其算法优势转化为挂钟加速。移动平均更新每个小批量都会更新每个集群的计数和总和。质心就是它们的比率。旧质心 新质心 批处理意味着 没有超参数调整 一切都源自缩放定律。 log n、log k、log T 开始和结束 新颖的初始化和终止方法。运行以每个集群的上一个纪元的确切平均值结束。开始精确均值余弦调度 步骤 α 从 α 0 退火到 0.1 α 0 ，就像学习率一样。 α₀ 0.1 α₀ 死簇复活 大小低于阈值的簇通过分裂大簇来复活。低于阈值的大集群分裂为两个专用内核 GPU 内核融合了最近质心分配，将上面的 epoch 节省转化为匹配的挂钟加速。朴素的 PyTorch 融合内核 质心更新规则 质心更新规则是 e-means 的核心组件。本节旨在概述 k 均值聚类的更新规则。所有方法都为每个簇 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 \] 因此较旧的分配会以速率 α 被遗忘，并且质心跟踪最近的分配而不是整个历史。该更新规则通常用于