用 Faiss 驱动的基于 ANN 的 KMeans

ANN-driven KMeans with Faiss

正如上一章所讨论的，KMeans 的训练过程如下：

• 步骤 1) 初始化质心
• 步骤 2) 为每个数据点找到最近的质心
• 步骤 3) 重新赋值质心
• 步骤 4) 重复直至收敛

但在这种实现中，“第 2 步”存在运行时瓶颈，因为这一步涉及暴力的穷举搜索。换句话说，它会计算每个数据点到每个质心的距离。

因此这一步并没有被优化，训练和预测都需要耗费大量时间。这在大型数据集上尤为棘手。为了加速 KMeans，我通常偏爱的一种实现——尤其是在大型数据集上——是 Facebook AI Research 的 Faiss。进一步说，Faiss 使用近似最近邻（ANN，approximate nearest neighbor）搜索算法，提供了快得多的最近邻搜索。它使用一种“倒排索引（Inverted Index）”，这是一种用于存储和索引数据点的优化数据结构。

这使得执行聚类变得极其高效，尤其是在大型数据集上，这一点从下图中也可以看出：

如上所示，在一个包含 50 万个数据点（1024 维）的数据集上，Faiss 大约比 Sklearn 的 KMeans 快 20 倍，这是一个极其惊人的加速。更重要的是，Faiss 还可以在 GPU 上运行，从而进一步提升聚类的运行时性能。可以从这里开始使用 Faiss：https://github.com/facebookresearch/faiss。