矢量搜索索引入门#

矢量搜索索引通常使用近似方法来权衡结果的准确性和速度,这可以通过降低延迟(端到端单个查询速度)或提高吞吐量(在短时间内可以满足的查询矢量数量)来实现。矢量搜索索引,特别是使用近似方法的索引,与机器学习模型非常密切相关,但它们针对快速搜索和结果准确性进行了优化。

当矢量数量非常少(例如少于 10 万个)时,使用暴力搜索(也称为平面索引)可能足够快,它返回精确结果,但代价是穷尽搜索所有可能的邻居。

目标#

本入门指南旨在解决矢量搜索索引的配置挑战,但其主要目标是帮助用户快速入门,并选择合适的索引类型和少量易于管理的调优参数,以获得足够满意的结果,而不是提供关于调优每个超参数的全面指南。

因此,我们重点关注 4 种主要的数据集大小

  1. 小型数据集(< 10 万个矢量),可能不需要 GPU

  2. 中型数据集(< 100 万个矢量),可能不需要 GPU

  3. 大型数据集(> 100 万个矢量),目标是快速创建索引,牺牲部分搜索质量

  4. 大型数据集,优先考虑高质量搜索,牺牲部分快速创建索引

与其他机器学习算法类似,矢量搜索索引通常有一个训练步骤(即构建索引)和一个推理或搜索步骤。超参数也倾向于分为构建参数和搜索参数。

虽然并非总是如此,但通常会观察到一种普遍趋势:搜索速度随着质量的提高而降低。索引构建性能也倾向于如此,尽管不同的算法在构建时间、质量和搜索时间之间存在不同的关系。重要的是要理解没有免费的午餐,因此每种索引类型总是存在权衡。

质量的定义#

当我们谈论索引的质量时,我们指的是什么?在机器学习术语中,我们使用召回率来衡量它,召回率有时与准确率互换使用,尽管这两者是略有不同的衡量标准。在矢量搜索中使用的召回率,本质上意味着“在我所有的结果中,哪些结果会包含在精确结果中?”在矢量搜索中,目标是找到与给定查询矢量最近的一些矢量,因此召回率往往比准确率更宽松,仅根据集合包含进行区分,而不是根据精确的有序列表匹配(这更接近于准确率衡量标准)进行区分。

选择矢量搜索索引#

许多矢量搜索算法通过将矢量空间划分为更小的部分来提高可扩展性,同时减少距离计算次数,这通常通过使用聚类、哈希、树和其他技术实现。另一种流行的技术是降低空间的宽度或维度,以降低计算每个距离的成本。

小型数据集(< 10 万个矢量)#

这些数据集非常小,使用 GPU 是否能带来任何价值值得怀疑。如果维度也相对较小(< 1024),您可以直接在 CPU 上使用暴力搜索或 HNSW,获得良好的性能。如果维度相对较大(1536, 2048, 4096),您应该考虑使用 HNSW。如果构建时间性能至关重要,您应该考虑使用 CAGRA 构建图并将其转换为 HNSW 图进行搜索(此功能目前存在于独立的 cuVS/RAFT 库中,并且很快将添加到 Milvus 中)。IVF 平面索引在这里也是一个不错的选择,因为它可以通过划分矢量空间从而减小搜索空间来提高搜索性能。

中型数据集,可能不需要 GPU(< 100 万个矢量)#

对于较低维度(例如 1024 或更小),您可以考虑在 GPU 上使用暴力搜索(即平面)索引,获得具有精确结果的非常好的搜索性能。您也可以在 CPU 上使用 HNSW 或在 GPU 上使用 CAGRA 等基于图的索引。如果构建时间至关重要,您甚至可以在 GPU 上构建 CAGRA 图并将其转换为 CPU 上的 HNSW 图。

对于更高维度(1536, 2048, 4096),在更高质量搜索设置下,HNSW 的构建时间性能会开始下降,因此构建 CAGRA 图会变得更有用。

大型数据集(> 100 万个矢量),目标是快速创建索引,牺牲部分搜索质量#

对于快速摄取且稍微降低搜索质量可接受(召回率 85% 及以上)的场景,IVF(倒排文件索引)方法非常有用,因为它们的构建速度非常快,并且仍具有可接受的搜索性能。IVF-flat 索引会将矢量划分为用户指定的簇数(n_lists),在搜索时,将对每个查询矢量使用暴力搜索对最接近的一些簇(由 n_probes 定义)进行搜索。

IVF-PQ 与 IVF-flat 类似,主要区别在于矢量使用有损乘积量化压缩,因此索引在 GPU 上的占用空间更小。通常建议将 n_lists 设置为 sqrt(n_vectors),并将 n_probes 设置为 n_lists 的一定百分比(例如 1%、2%、4%、8%、16%)。由于 IVF-PQ 是有损压缩,可以通过初步增加邻居数量(通过某个倍数因子)并使用原始矢量计算精确距离来执行精炼步骤,最终将邻域缩小到 k 大小。即使是 2 倍的精炼(最初查询 k*2 个)也能非常有效地弥补 PQ 压缩造成的召回率损失,但这需要保留原始矢量(许多数据库无论如何都会存储原始矢量)。

大型数据集(> 100 万个矢量),优先考虑高质量搜索,牺牲部分快速创建索引#

通过牺牲索引创建性能,可以构建极高质量的搜索模型。通常,所有矢量搜索索引类型都具有与搜索准确率直接相关的超参数,因此可以提高这些参数以获得更好的召回率。不幸的是,这也会显著增加索引构建时间并降低搜索吞吐量。这里的关键是找到最快的构建时间,同时实现最佳召回率、最低延迟或最高吞吐量。

关于推荐的索引类型,像 HNSW 和 CAGRA 这样的基于图的算法在大数据集上通常具有很好的扩展性,同时在质量方面具有卓越的搜索性能。挑战在于基于图的索引需要学习一个图,因此,正如本节副标题所示,它们的构建速度通常比其他选项慢。在 GPU 上使用 CAGRA 算法可以显著缩短构建时间,并且与在 CPU 上搜索相比,具有更高的吞吐量(和更低的延迟)。目前,在 GPU 上使用 CAGRA 的缺点是需要图和原始矢量都适合 GPU 内存。一种折衷方案是在 GPU 上构建 CAGRA 图,然后将其转换为 HNSW 以便在 CPU 上进行高质量(和中等速度)的搜索。

调优和超参数优化#

不幸的是,对于大型数据集,对整个数据集进行超参数优化并不总是可行。但是,可以在较小的子集上进行超参数优化,并找到应该能很好地泛化到整个数据集的合理可接受的参数。通常,这种超参数优化需要使用暴力搜索等精确方法在子集上计算一个真实结果,然后用它来评估在随机抽样矢量上的多次搜索。

完全的超参数优化也并非总是必要的——例如,一旦您在子集上构建了真实结果数据集,很多时候您可以从使用默认构建参数构建索引开始,然后尝试不同的搜索参数,直到获得所需的质量和搜索性能。对于可能达到数 TB 的大型索引,您也可以对例如 1000 万个矢量进行子采样,训练一个索引,然后从那里调优搜索参数。尽管可能存在少量误差,但对于构建本地分区索引的数据库来说,选择的构建/搜索参数应该能很好地泛化。

矢量搜索索引类型总结#

名称

权衡

最适合用于…

暴力搜索(即平面)

精确搜索,但需要穷尽距离计算

小型数据集(< 10 万个矢量)

IVF-Flat

划分矢量空间以减少暴力搜索的距离计算,但以牺牲召回率为代价

中型数据集(< 100 万个矢量)或大型数据集(> 100 万个矢量),优先考虑快速构建索引而不是质量。

IVF-PQ

在 IVF-Flat 中添加乘积量化,以牺牲召回率为代价实现规模扩展

大型数据集(远大于 100 万个矢量),优先考虑快速构建索引而不是质量

HNSW

显著减少距离计算,但以更长的构建时间为代价

中型数据集(< 100 万个矢量)或大型数据集(> 100 万个矢量),优先考虑搜索质量和速度而不是索引构建时间

CAGRA

显著减少距离计算,但以更长的构建时间为代价(尽管构建时间比 HNSW 短)

大型数据集(远大于 100 万个矢量),优先考虑搜索质量和速度,但索引构建时间也仍然重要。

CAGRA 构建 + HNSW 搜索

(即将添加到 Milvus)

显著减少距离计算并缩短构建时间,但以更高的搜索延迟 / 更低的吞吐量为代价。大型数据集(远大于 100 万个矢量),索引构建时间和搜索质量重要,但 GPU 资源有限且搜索延迟不敏感。