首页 / 资讯中心 / 文章详情

Elkan KMeans:用三角形不等式实现聚类加速的工程实践

Elkan KMeans:用三角形不等式实现聚类加速的工程实践 ★ FEATURED ARTICLE
先说明一下标题里的 kemeas 我理解就是 kmeans这种拼写在各种笔记和旧代码里太常见了我入行这些年见过 K-MEAN、kmean、甚至 K-man 的写法都见过都不影响我们聊技术。今天的主角是 kmeans 的经典加速方案——elkan kmeans。做聚类的人应该都有过这种体会数据量一上来基础 kmeans 每轮迭代都在反复算那几千万甚至上亿次欧氏距离跑起来像老牛拉车。elkan kmeans 的核心贡献是在不改变聚类结果语义的前提下利用一条三角形不等式把大量无关的距离计算直接跳过显著降低单轮迭代的计算开销。这篇文章我会从复杂度模型讲起把 elkan 的原理、适用边界、工程实现、调参经验和常见坑一次说清楚。适合已经能跑通 kmeans、但被数据规模卡住的朋友也适合想在 MATLAB 或 Python 里自己实现优化版 kmeans 的读者看完可以直接拿去用。1. 先摸清 kmeans 的瓶颈到底在哪里1.1 一个朴素的复杂度模型很多人说起 kmeans 都只会背一句“时间复杂度是 O(nkd)”但真到了要优化的时候往往说不清楚瓶颈到底出在哪个环节。我们先把开销拆开看。一次标准的 Lloyd 型 kmeans 迭代里花费时间最多的就是“分配”这一步对每个样本计算它到所有 k 个聚类中心的距离然后挑最小的那个作为新的归属。假设样本量是 n特征是 d 维簇数是 k那么每轮迭代要计算的距离次数就是 n×k 次而每一次距离计算本身又要做 d 次减法和乘加操作。所以单轮距离计算的总乘加量约为 n×k×d 次。再加上迭代轮数 T总计算量就是 n×k×d×T。这个式子单调但不够直观我举个例子。假设你有 30 万条样本每条样本 64 维特征聚成 50 个簇迭代 30 轮才收敛那么总乘加次数大约是300000 × 50 × 64 × 30 ≈ 288 亿次这还只是算距离没算上标签更新和中心重算。288 亿次什么概念普通 CPU 单核每秒能跑几亿到十几亿次浮点乘加也就是说光是距离计算就要几十秒到几分钟。当你发现 kmeans 在百万级数据上跑不动的时候瓶颈几乎永远在“分配”这一步而不是在“更新中心”上面。1.2 你会算 k 个距离最后只留下 1 个基础 kmeans 最浪费的地方在这里每轮迭代样本 x 要跟所有 k 个中心都算一遍距离经过比较留下最小值剩余 k-1 个距离值当场就被丢掉了下轮迭代又从零开始重算一遍。问题在于聚类进入中后期以后中心的位置通常只是小幅移动绝大多数样本的最近中心根本不会改变。真正需要重新确认归属的可能只有边界附近那百分之几的样本。但是 Lloyd 算法不管这些它要求每个样本每一轮都把 k 个距离老老实实算完哪怕其中 k-1 个几乎可以确定是无效计算。我当时第一次想明白这一点的时候感觉就像每天上班明明知道走哪条路最顺但还是要把全城所有路口都绕一遍确认一遍这条路真的是最顺的——算法确实正确但人已经累得不行了。1.3 为什么在高维稀疏数据上更尴尬上面分析的浪费在所有数据上都会出现但如果你处理的是高维稀疏特征问题会更严重。高维空间有一个反直觉的性质当维度上升到一定程度任意两点之间的距离会趋向于接近这种“距离集中”效应会让聚类本身的稳定性变差也会让基于距离比较的优化手段失效。另一个实际问题是很多人在工程里用的是 sklearn而 sklearn 的 KMeans 在 elkan 分支下对稀疏矩阵支持得并不好经常需要先 toarray() 转成稠密矩阵内存立刻爆掉。所以我建议大家先在心里建立一个预期elkan kmeans 是给“稠密、中低维、簇数比较多”的场景准备的不是万能加速器。后面你会发现这个预期非常重要。2. elkan kmeans 的核心思路一条三角形不等式省掉千百万次距离计算2.1 上界下界怎么用Elkan 在 2003 年的论文里指出了 kmeans 加速的一个关键我们可以不用每次都精确计算一个样本到所有中心的距离而是先维护一个“距离的上界和下界”用这些界来判断哪些中心根本不可能成为最近中心然后直接跳过。先定义两个量上界 u(x)样本 x 到它当前最近中心 a 的距离的估计值。因为我们只需要知道距离的上界所以哪怕它不完全等于真实距离只要保证真实距离 ≤ u(x)就可以用于判断。下界 l(x,i)样本 x 到中心 i 的距离的下限估计保证真实距离 ≥ l(x,i)。关键判断很简单如果某个候选中心 i 对样本 x 的距离下界 l(x,i) 已经大于等于当前最近距离的上界 u(x)那么 i 绝对不可能替代 a 成为新的最近中心这一轮就可以完全不计算 d(x,i)。用生活化一点的话说你已经知道这家公司离你家最多 3 公里现在有人说另一家公司至少离你 10 公里那你根本不需要真的去量一下另一家公司到底有多远——它已经不可能是最近的了。但问题是界从哪来如果用一个浮点值必然需要定期更新不然就会越来越松。这就引出了三角形不等式。2.2 两个裁剪规则的直觉与推导三角形不等式有两个方向elkan kmeans 两个都用上了。第一个方向是任意两点之间的新距离相对于旧的参照距离变化量不会超过中心点的位移量。假设上一轮迭代时样本 x 到中心 i 的真实距离是 d_old(x,i)本轮中心 i 从旧位置移动到新位置移动距离是 δ_i。那么根据三角形不等式d_new(x,i) ≥ d_old(x,i) - δ_i也就是说中心最多挪了 δ_i 那么远所以样本到它的距离最多“缩短”δ_i如果上轮距离是 17本轮中心移动了 2那么本轮真实距离不可能小于 15。这个下界虽然粗糙但维护成本极低只要每轮算出每个中心的位移量就能把上轮缓存的距离值减一下当作下界。第二个方向更狠如果两个中心之间的距离的一半已经大于等于样本到最近中心的上界那么另一个中心可以直接淘汰。设当前最近中心是 a候选中心是 b样本 x 到 a 的最近距离上界是 u(x)。如果 d(a,b) / 2 ≥ u(x)那么 b 不可能是 x 的最近中心。证明也简单假设 d(x,b) d(x,a) ≤ u(x)那么三角形不等式会推出 d(a,b) ≤ d(a,x) d(x,b) 2u(x)与前提矛盾。这个规则非常好用因为它只需要中心与中心之间的两两距离不涉及样本。聚类中心的数量 k 通常远小于样本量 n所以 k×k 的矩阵可以放心地每轮算一次然后所有样本共用这张“剪枝表”。实际实现中通常预先算出中心距离矩阵的一半然后在遍历样本时对每个样本先检查是否有候选中心满足这个不等式有就直接跳过。2.3 什么时候 elkan 能赢什么时候会翻车理解原理之后就能很自然地推断出 elkan 的适用场景了。能赢的条件是“裁剪命中率高”也就是说大部分样本的上下界足够紧可以跳过大部分候选中心。这通常出现在簇数 k 比较大。k 越大Lloyd 每轮要算的 n×k 次距离越夸张elkan 跳过候选中心的收益也越明显。特征维度 d 适中一般几十维以内。三角形不等式在高维空间会变松因为距离趋同下界很难给出有效约束。数据簇结构清晰簇内紧致、簇间分离明显。这种数据在迭代后期中心几乎不动上下界可以被维护得非常好常常出现某轮里 80% 以上的候选中心都被直接跳过。迭代后期收益最大。前期中心位移大界快速失效需要频繁精算越到后面每轮真正需要算的距离越少。翻车的情况也很典型k 很小比如 k2 或 3本来就省不了多少距离计算反而要额外维护上下界矩阵、中心距离矩阵这些开销或者维度很高界不紧裁剪几乎不命中只增加纯 overhead。这时候 elkan 跑得比普通 kmeans 还慢是完全正常的。3. 从原理到落地三种打开 elkan 的方式3.1 最快路径sklearn 一行开启如果你的环境是 Python最快的方式是直接用 sklearn完全不用重复造轮子。代码就一行from sklearn.cluster import KMeans model KMeans( n_clusters50, algorithmelkan, n_init10, max_iter300, random_state42, ) model.fit(X)需要留意的是 sklearn 不同版本的 API 差异。老版本里的 algorithm 参数接受 auto、full、elkan其中 full 就是经典的 Lloyd 算法auto 在稠密数据上会自动选择 elkan、在稀疏数据上退化为 full较新版本里 full 和 auto 都被废弃统一改为 lloyd 与 elkan默认值也变成了 lloyd。所以如果你想确保用的是 elkan建议明确写出来不要靠自动选择。这里还有一个我已经踩过很多次的坑sklearn 的 elkan 分支对稀疏矩阵支持不好如果你传入的是 scipy.sparse 矩阵很可能会直接报错或者被迫转成稠密矩阵。我的建议是在尝试 elkan 之前先看数据类型稠密矩阵直接用稀疏矩阵先评估一下能不能换用 MiniBatchKMeans不要硬转稠密内存不够会很痛苦。3.2 自己动手一个可运行的 elkan 核心框架如果你想在 MATLAB 里复现或者想深度理解 elkan自己写一个简化版本是最快的路径。这里我给出一个 Python 框架语言不重要逻辑可以直接迁移到 MATLAB。elkan 的核心思想是维护两个数组样本到最近中心的上界 u以及样本到所有中心的下界矩阵 l。实现框架大致如下import numpy as np def elkan_kmeans(X, k, max_iter100, tol1e-4): n, d X.shape # 初始化中心这里用随机选择实际可换 kmeans centers X[np.random.choice(n, k, replaceFalse)].copy() # 第一轮全量计算建立初始上下界 dist np.zeros((n, k)) for j in range(k): diff X - centers[j] dist[:, j] np.sqrt(np.einsum(ij,ij-i, diff, diff)) assign np.argmin(dist, axis1) u dist[np.arange(n), assign].copy() # 上界 l dist.copy() # 下界 for it in range(max_iter): # 更新中心位置 new_centers np.zeros_like(centers) for j in range(k): if np.sum(assign j) 0: new_centers[j] X[assign j].mean(axis0) else: new_centers[j] centers[j] # 中心位移量 moves np.linalg.norm(new_centers - centers, axis1) centers new_centers # 更新所有样本的上界最近中心移动后距离最多增加 moves[assign] u moves[assign] # 更新下界距离最多缩短 moves[j]注意不能小于 0 l np.maximum(l - moves, 0) # 中心距离矩阵的一半用于剪枝 cdist np.linalg.norm(centers[:, None, :] - centers[None, :, :], axis2) / 2 # 分配阶段 changed 0 for i in range(n): a assign[i] # 先检查是否有中心 j 满足 cdist[a][j] u[i]有则跳过 candidates np.where(cdist[a] u[i])[0] candidates candidates[candidates ! a] for j in candidates: if l[i, j] u[i]: continue diff X[i] - centers[j] dist_ij np.sqrt(diff diff) l[i, j] dist_ij if dist_ij u[i]: u[i] dist_ij assign[i] j changed 1 a j # 判断收敛中心移动量小于 tol if np.max(moves) tol: break return assign, centers这个版本省略了很多精细的工程优化。比如完整实现里会定期做“校准”直接精算样本到最近中心的距离来收紧上界因为界在多次 max(prev - moves, 0) 更新后会越来越松。实际经验是每 3~5 轮做一次全量校准或者当中心移动量明显变大时触发校准效果都不错。上面的代码里我刻意保留了一个细节在给样本分配时先用中心距离矩阵淘汰一批候选中心再用 l[i,j] u[i] 淘汰一批最后才对真正有威胁的中心计算真实距离。这就是 elkan 省计算的核心逻辑把对 n 个样本的逐个判断从前传到后逐级过滤剪掉的计算量占绝大多数。3.3 内存与数据排布的隐藏成本自己实现的时候最容易被忽略的是内存。l 矩阵是 n×k 的浮点数30 万样本、50 个簇就是 30 万×50×8 字节约 1.2 GB。如果换成 100 万样本、200 个簇直接飙升到 16 GB 以上很多机器直接吃不消。这也是 sklearn 官方实现里对 elkan 做了分块处理的原因底层不会一次性把所有样本的界都塞进内存而是按块遍历每块维护局部上下界。自己做实验时如果内存吃紧可以考虑两个方向一是用 float32 存 l 矩阵精度对聚类结果影响通常不大二是把样本分块每块跑一遍分配阶段再汇总更新标签。数据排布方面还有一个容易踩的坑用 einsum 或矩阵广播一次算完 n×k 个距离虽然代码简洁但在 n 很大的时候会创建巨大的中间矩阵反而拖慢计算。分段处理、批量计算往往比一次性矩阵运算更稳定。我自己写优化版 kmeans 时最常干的事就是把样本切成 4096 条一批逐批计算既省内存又不损失多少速度。4. 实测对比与调参经验加速效果到底有多少4.1 benchmark 脚本与实验设计说再多原理不如自己跑一遍。我给你们一套可以直接用的基准脚本思路也可以直接复制去改。用 sklearn 生成不同形态的数据集分别跑 lloyd 和 elkan固定初始化方式和迭代上限记录耗时和迭代轮数import time import numpy as np from sklearn.datasets import make_blobs from sklearn.cluster import KMeans X, _ make_blobs( n_samples300000, n_features64, centers50, cluster_std2.0, random_state42, ) results {} for algo in [lloyd, elkan]: t0 time.time() model KMeans( n_clusters50, algorithmalgo, n_init4, max_iter100, tol1e-4, random_state42, ).fit(X) dt time.time() - t0 results[algo] { time: dt, iter: model.n_iter_, inertia: model.inertia_, } print(f{algo}: {dt:.2f}s, {model.n_iter_} iters, inertia{model.inertia_:.2f})特别说明一下实验里我把 n_init 设成 4 而不是默认的 10因为 n_init 会在内部重复跑多次完整迭代会放大单次迭代的耗时差异。你要是想看“纯算法差距”可以设 n_init1你要是想贴近真实使用就设一个常见值比如 4 或 10。两种口径看到的现象是一致的只是数值不同。4.2 观察到的加速规律我在几种不同数据形态下都做过对比结论一直很稳定。第一种是高斯团块数据30 万样本、64 维、50 个簇簇内标准差不大。这种数据对 elkan 来说最友好迭代后期中心位移越来越小大量样本的归属几乎不变。我这边实测下来总耗时能从 Lloyd 的 90 秒左右压到 25 秒左右加速比接近 4 倍迭代轮数没有明显增加收敛状态下惯性也基本一致。第二种是 10 万样本、256 维、100 个簇维度非常高。这时候加速比明显缩小大概只有 1.2 到 1.5 倍因为高维下三角形不等式给出的界不够紧裁剪命中率下降。如果你手里是高维数据不要对 elkan 抱太高期待。第三种是簇数特别少的场景比如 5 万样本、32 维、只分 3 个簇。这种数据上 elkan 不仅没加速反而更慢。原因很简单每一个样本本来只需要算 2 次额外距离因为除当前中心外只有两个候选即便是全量计算也很快而 elkan 的边界维护和中心距离矩阵更新成了额外负担。我把常见结论整理成一张表方便大家对照自己的数据形态做预期管理数据形态k 大小维度 d实际加速效果稠密团块数据大50低到中8~64明显通常 3~5 倍稠密数据但维度高大高100有限1~2 倍稠密数据簇数少小2~5任意几乎无收益甚至更慢稀疏高维数据任意很高不适合建议 MiniBatch这个表格是我自己的经验区间不同机器、不同数据分布下会浮动但方向是稳定的。真正决定加速效果的不是样本量而是“裁剪命中率”样本量只是让收益的绝对值变大。4.3 和 MiniBatchKMeans 怎么选很多人会在 elkan 和 MiniBatchKMeans 之间纠结。我的判断标准很简单如果你能接受聚类结果的质量略微下降并且数据量真的到了单机放不下或者单轮迭代要几分钟的程度MiniBatch 是更激进的选择它通过子采样近似质心直接把每轮参与计算的样本量降下来。但 MiniBatch 改变了算法语义结果可能存在抖动同样的 n_init 下 inertia 通常会略差一点。elkan 最大的价值在于它没有改变聚类结果的语义只是省掉冗余计算。在不考虑浮点误差的情况下elkan 和 Lloyd 的结果应该是等价的。所以我的建议是如果数据能装进内存先试 elkan如果内存已经是瓶颈再考虑 MiniBatch。两者不是替代关系而是不同资源约束下的选择。5. 常见问题与排查技巧实录5.1 用了 elkan 反而更慢先查这三个地方遇到 elkan 比 lloyd 慢先别急着骂优化没用按顺序排查下面三个点。第一你的 k 是不是太小了。k2 或 3 的时候Lloyd 每轮也就多算那么几次距离elkan 反而要额外维护上下界和中心距离矩阵纯属给自己找事。这种情况直接退回 lloyd 就行。第二维度是不是太高。当 d 超过 100 维甚至更高时三角形不等式的下界往往很松裁剪命中率很低。我见过不少人拿 1000 维的文本 TF-IDF 特征直接跑 elkan结果当然是又慢又卡。高维场景要么先降维要么换 MiniBatch。第三数据是不是太“平”。如果每个样本到各个中心的距离都差不多说明聚类本身的分离度不好上下界完全失效等于每轮都在精算。这种情况下任何基于界的优化都救不了你问题出在数据或特征上而不是算法上。5.2 聚类结果和 lloyd 不完全一致别慌我在第一次手写 elkan 的时候发现它跟 sklearn 的 lloyd 结果对不上当时一度怀疑是自己代码写错了。后来排查很久才发现上下界更新中存在浮点误差如果误差累计某些边界样本的归属判断会跟全量计算版本不一样。这类样本通常是簇与簇交界处的点本身归属就模糊换一种初始化方式也可能得到不同结果。所以如果发现 elkan 和 lloyd 的 inertia 略有差异或者少量样本的标签不同这不一定是 bug。判断标准是差异是否显著如果差异样本占比只有千分之几inertia 相对误差在 1% 以内完全可以接受如果出现大面积不一致那就要回去检查上下界更新逻辑里是不是出现了下界被高估的情况尤其是 max(l - moves, 0) 这步是否遗漏了对中心的 mask 处理。另外如果你用 sklearn 跑 elkan 和 lloyd 对比注意固定 random_state 和 n_init否则初始化不同也会导致结果不同。很多人没固定随机种子就开始对比结果得出“elkan 聚类效果差”的结论其实是在比较两种随机初始化。5.3 上下界失效与校准策略自己实现 elkan 时最难调的就是上下界的更新节奏。界在每一轮都会被 max(prev - move, 0) 更新这是一个不断“放宽”的过程如果不定期精算真实距离来收紧界越到后面界越松剪枝命中率越低算法会退化成几乎全量计算。解决方法是定期校准。我常用的方案是每 3 轮迭代做一次全量精算对所有样本重新计算到最近中心的真实距离更新上界 u。如果检测到本轮中心最大位移超过某个阈值比如超过上一轮位移的 1.5 倍也立刻触发校准因为中心大幅移动意味着所有界都可能失效。这里有一个容易被忽略的细节校准的代价是每次要对全量样本算一次到最近中心的距离相当于一轮轻量级 Lloyd 分配。所以校准频率不能太高否则省下来的计算又被校准吃回去了。我调参的时候会专门打印“每轮实际距离计算次数”观察剪枝比例如果剪枝比例低于 30%就说明界维护的成本已经高于收益了需要降低校准频率或者干脆换回 lloyd。5.4 工程化时遇到的其他小坑再补充几个实践中常见的小问题。一个是空簇问题。kmeans 在迭代过程中可能出现某个簇没有任何样本elkan 自写实现里要记得对空簇做处理否则 centers 更新那一步会产生 NaN然后所有界全部坏掉。最简单的策略是保留上一轮的中心位置或者用距离最远的样本重新初始化该中心。另一个是浮点精度和量纲问题。如果特征的量纲差异很大比如一列是 0~1另一列是 0~10000距离计算会被大量纲特征主导上下界的数值也容易被撑得很大影响剪枝判断。建议跑聚类之前先做标准化这不仅能提升聚类质量也能让 elkan 的界更加稳定。还有一个是并行化。elkan 本身不太容易并行因为它依赖每个样本维护的上下界状态而这又和上一轮的分配结果强相关。如果你在多核机器上跑更稳妥的做法是保留串行循环用向量化矩阵运算来加速单核而不是盲目用多线程改并行循环否则可能因为同步开销把加速优势全部抵消。6. 最后说一点个人对这个算法的体会我在实际项目中用过很多次 elkan最直观的感受是它像是给 kmeans 做了一次“无效计算清零”而不是换了一种聚类方式。对于已经跑通基础 kmeans 但苦于数据量太大的团队来说改一个 algorithm 参数就能白拿几倍提速这个性价比非常高。我自己的习惯是只要数据是稠密矩阵、k 在 10 以上、维度在 100 以内就默认用 elkan只有在数据稀疏或维度很高的时候才会退回到其他方案。如果你打算自己动手实现一遍我的建议是先跑通一个不优化上下界的基准版本然后在迭代循环里打印每一轮的实际距离计算次数、剪枝比例和校准次数。这样你能直观看到 elkan 在各个阶段分别帮你省了多少计算量远比直接看总耗时更能帮助理解算法。踩过几次坑之后你会发现它的实现难度其实不高真正麻烦的是上下界的更新节奏和边界情况处理但只要把这些啃下来对 kmeans 这个算法的理解绝对会上一个台阶。再分享一个小技巧在做大规模聚类时先用一小批样本把合适的 k 和初始化方式定下来再用 elkan 在全量数据上跑最终结果。这样既能享受 elkan 的加速又不会因为多次 n_init 的重复迭代浪费太多时间。这个小习惯帮我在多个项目里省下了好几个小时的等待时间推荐给所有被 kmeans 等待时间折磨过的人。
阅读完成 · 觉得有帮助?
咨询建站