AI 应用开发:你知道怎么用好 pgvector 吗(下篇)
任何数据库,要真正发挥其能力,最重要的技能是学会使用索引,没有之一。
如果你还没有看过第一篇文章,我建议你先阅读它,链接放在文章开头了。上篇详细探讨了 pgvector 的重要概念和 AI 应用,并提供了一个实际例子,展示了如何基于词义而非词本身进行搜索。
在这篇文章中,我们将探讨 pgvector 支持的索引的更多细节。我们将讨论这些索引在后端是如何构建的,以及与这些索引相关的各种参数,并指导你根据你的需求选择最合适的索引。
最后,我们将评估在我们的一百万个来自维基百科的记录数据集中,哪种索引为我们的搜索查询提供了最佳的召回率。让我们深入了解 pgvector 支持的两种索引类型:Inverted File(IVFFLAT)和 Hierarchical Navigable Small World(HNSW),它们都是近似最近邻(ANN)索引,意味着它们可以高效地搜索与查询向量足够接近的邻居向量。
请注意:为了在完整列上创建索引,该列的所有行必须具有一致的向量维度数量。例如,如果一个名为 country 的表有一个名为 city 的向量列,而这一列中的一行包含 1536 个维度的向量数据,而下一行只包含 500 个维度,我们将无法在该列上创建完整索引。如果维度变化,考虑创建部分索引。
IVFFLAT
这是 pgvector 中首先引入的索引。当我们生成 IVFFLAT 索引时,它会找到看起来相似的向量并形成多个簇,然后为每个簇计算一个中心点,称为质心。然后,每个向量被分配到其最近的质心。这种聚类确保在相似性搜索期间,我们只寻找与查询向量接近的向量,排除不那么相似的向量。这可以提高搜索速度。下面的图表显示了如何形成质心以及如何将相似的向量分配到一个区域。
IVFFLAT 的构建
当我们生成 IVFFLAT 索引时,我们只能在我们的表中有足够的数据时才能创建。这是由于几个原因。我们使用的聚类算法依赖于有足够的数据点来准确捕捉底层数据分布并形成有意义的簇。每个簇的质心,即每个簇的代表点,是基于簇内的数据点计算的。将质心映射到每个簇内的向量的倒排文件在簇准确反映数据相似性时更有效。虽然没有足够的数据点的严格阈值,但通常建议在构建索引之前至少有几千个数据点,以确保有意义的聚类、准确的质心和有效的性能。
以下是创建索引的查询示例:
在上面的查询中,列表(lists)决定了我们在构建索引时想要创建多少个簇。更多的列表可以通过快速缩小潜在匹配项来加快搜索速度。然而,太多的列表可能会对召回率产生负面影响,因为探索较少的子组可能会错过相关的邻居。选择列表的推荐值是行数除以 1000,对于最多 1 百万行,即如果表中有 100 万行,则列表的值应为 1000。对于超过 1 百万行,即如果表中有 200 万行,则列表的值应为 1415。相似性搜索方法可以是 Vector_l2_ops、Vector_ip_ops 或 Vector_cosine_ops。不熟悉这些术语?请查看我们之前的博客以获取详细解释。
构建索引时的注意事项
我们还需要在构建索引时考虑 maintenance_work_mem,因为构建索引涉及内存密集型操作,如距离计算和数据比较。因此,在开始创建索引之前,最好分配足够多的内存。没有关于良好起点值的规则或建议,但可以作为一个起点:对于较小的数据集(最多几百万个向量),512MB 到 1GB 通常就足够了。对于较大的数据集(数千万个或更多),可能需要 2GB 到 4GB 甚至更多的内存。我们还可以考虑在索引构建期间临时增加 maintenance_work_mem,然后在完成后减小它。或者,可以在会话级别设置它。pgvector 提供了关于有效创建索引所需的内存的提示。例如,尝试在默认 64MB 的 maintenance_work_mem 下对 100 万行创建索引时,我遇到了以下错误:
IVFFLAT 的潜在性能问题
IVFFLAT 是一种近似最近邻(ANN)算法,它专注于速度而不是在识别最近邻居时的精确性。在搜索过程中,它首先识别与查询点最近的簇质心。然而,这种方法可能会忽略那些位于这些簇之外或比其各自簇质心更接近查询点的相关向量。因此,尽管 IVFFLAT 提供了快速的搜索能力,但在可能错过一些确切相似向量方面存在权衡。
探针(Probes)
探针解决了这个挑战,它允许 IVFFLAT 探索更广泛的簇或区域,提高了发现实际最近邻居的机会。通过指定额外的探针,我们可以在更多区域进行搜索,提高搜索质量和减少忽略最接近向量的风险。需要注意的是,虽然使用更多的探针可以通过检查额外的簇来提高准确性,但它也会带来搜索时间增加的权衡。探针的推荐值是列表的平方根,即如果你有 1000 个列表,那么探针应该是 31。默认情况下,其值为 1,这意味着它只会在最近的质心簇/区域中搜索,而不会考虑其他区域。
重建索引
IVFFLAT 的结构严重依赖于初始数据集分布。它最初计算聚类、列表、质心和数据点。但在现实世界中,当数据被添加、更新或删除时,索引可能会变得碎片化,导致搜索速度变慢、内存使用增加和结果不够准确。因此,重建索引允许它根据更新后的数据重新聚类和重新生成倒排文件,优化搜索准确性和性能。在数据分布发生重大变化后,重建索引总是一个好主意。如果搜索速度或准确性明显下降,可能需要更频繁地重建。对于大型数据集,即使没有明显的问题,也可以考虑定期重建。
通过 IVFFLAT 索引进行搜索
当搜索向量出现时,我们首先计算查询向量与所有质心之间的距离。在识别出最近的质心后,我们接着评估查询向量与该选定质心链接的每个向量之间的距离。最终,我们选择距离最近的向量作为结果向量。这样可以确保选择最相似的向量作为结果。在以下情况下考虑使用 IVFFLAT 索引:
可以容忍搜索准确性略有下降。
没有严格的内存限制。
数据更新不频繁。
如果希望索引构建快速且占用空间较小。
HNSW 索引
在 pgvector 0.5.0 版本中推出的这种索引类型以其显著特点脱颖而出。它在查询性能上优于 IVFFLAT 索引,但构建速度较慢,占用更多内存。与 IVFFLAT 不同,创建 HNSW 索引不需要训练步骤。即使表是空的,你也可以在设置关系/表后立即创建它。
HNSW 索引使用 SkipList 和 Navigable Small World Graphs 的概念,在高维空间中实现高效的 ANN 搜索。SkipList 是一种概率性数据结构,旨在通过引入多级访问(或“跳过”)来提高传统链表的搜索性能。想象一下翻阅一本大型电话簿——传统的链表需要按顺序阅读每个名字,而 SkipList 在特定间隔提供捷径以加速搜索。列表中的每个元素都是一个节点,包含一个值和指向其他节点的链接。不同级别的节点通过额外的“跳过”指针垂直连接。底层是一个传统的单维链表,其中所有节点按顺序连接。更高级别的节点较少,每个节点跳过下一层的一定数量的节点。
如何在 SkipList 中进行搜索
我们在 SkipList 的最高级别开始搜索,将目标元素与当前级别的下一个元素进行比较。如果下一个元素小于目标元素,向前移动到同一级别的下一个元素。如果下一个元素大于目标元素,向下移动到下一层并重复第二步。继续向前或向下移动,直到找到目标元素或到达 SkipList 的末尾。如果找到目标元素,返回元素或其位置。如果找不到目标元素,返回一个指示元素不在 SkipList 中的信号。
Navigable Small World Graph (NSWG)
NSWG 是一个图结构,其中每个顶点都连接到几个其他顶点,形成一个好友列表。在 NSWG 中进行搜索时,我们从一个与附近顶点相连的入口点开始,通过贪婪路由迭代地移动到最近的顶点。这个过程一直持续到我们达到一个局部最小值,表示没有更近的顶点可用。
HNSW 结合了 SkipList 和 NSWG 的概念,创建了一个强大的算法,用于在高维空间中进行高效的最近邻居搜索。它的核心是一个 Navigable Small World Graph 结构,确保网络中的节点保持局部聚类和短路径长度。这意味着相似的数据点紧密连接,便于在局部邻域内快速搜索。每一层代表数据的不同粒度级别,允许在搜索过程中更有效地导航。与 SkipList 中的层次结构类似,这些层次结构提供了捷径,使算法能够跳过不必要的比较,并更快地收敛到最近邻居。
在下面的图表中,我们可以看到形成一个倒排树的层次结构。最上层由一组精心选择的数据点组成,通常称为入口点。随着我们向下穿过层次,数据点的数量增加,底层包含数据集中的所有数据点。每一层不仅垂直连接到上层或下层,而且水平连接到同一层,创建了一个多层图结构。这种连接网络使得数据集的导航和搜索变得高效,并允许算法在局部和全局搜索之间创建平衡,以找到最近邻居。
HNSW 索引的创建
要创建 HNSW 索引,请使用以下查询:
参数 M
在 HNSW 图中,参数 M 指定每个顶点/节点与其最近邻居的最大连接数。它影响图的连通性,较高的值表示更密集的连接。较高的 M 允许更快的图探索,可能导致更快的搜索。然而,它也可能增加错过更远的相关邻居的可能性。相反,较低的 M 减少了探索选项,可能会减慢搜索速度,特别是在到达远距离邻居时。较低的 M 值可能通过迫使搜索探索更广泛的范围来提高准确性(召回率),增加找到相关邻居的机会。
ef_construction
这个参数定义了算法在将新数据点添加到图时考虑的最大候选邻居数。较高的值在构建过程中考虑更多的潜在相关邻居,可能导致更准确、连接更好的图。然而,评估更多候选项需要更多时间,导致索引创建时间更长。较低的值可以加快索引构建时间,因为需要较少的评估工作。然而,这可能导致索引质量较低,因为探索的候选项较少,可能影响搜索性能。
ef_search
这个参数定义了 HNSW 索引搜索操作期间搜索算法将考虑的最大候选邻居数。较高的值可能导致搜索速度变慢,因为需要评估更多的候选项,进行更多的比较和搜索时间。然而,它们提供了更高的召回率,通过探索更多的邻居,增加了找到相关邻居的机会,即使它们不是绝对最近的。相反,较低的值可以通过评估较少的候选项来加快搜索速度,可能导致更快的搜索时间。然而,它们也可能由于探索较少而导致召回率降低,可能会错过一些相关邻居,影响召回率。
何时使用 HNSW 索引
考虑在以下情况下使用 HNSW 索引:
当你需要高召回率速度时。
如果数据正在快速变化。
如果你想要存储高维数据。
如果可以容忍召回率的准确性。
测试索引的召回率和速度
现在让我们深入到示例中,看看通过在 pgvector 中使用索引可以获得什么样的性能提升。我们创建了 3 个表:没有索引的表(demo_table)、带有 IVFFLAT 索引的表(wiki_data_ivf)和带有 HNSW 索引的表(wiki_data)。我们使用 pgvector 版本 0.6.0 和 PostgreSQL 版本 16.1,并用来自 dbpedia-entities-openai-1M 的相同数据集填充它们。我们对一个查询运行 EXPLAIN ANALYZE,该查询通过选择与“<EMB>”的余弦距离超过 0.8 的文本,并仅选择 1 行。
没有索引的召回
并行序列扫描在 demo_table 上(成本=0.00..61781.66 行=138867 宽度=339)(实际时间=33.277..4119.286 行=41 循环=3)。规划时间:1.982 毫秒。执行时间:4145.951 毫秒。
带有 IVFFLAT 索引的召回
使用 wiki_table_ivf_openai_idx 索引扫描 wiki_table_ivf(成本=2512.50..7327.33 行=333333 宽度=340)(实际时间=12.434..12.434 行=1 循环=1)。规划时间:1.005 毫秒。执行时间:12.491 毫秒。
带有 HNSW 索引的召回
使用 wiki_table_openai_idx 索引扫描 wiki_table(成本=384.72..2063911.53 行=333433 宽度=340)(实际时间=8.548..8.549 行=1 循环=1)。规划时间:0.937 毫秒。执行时间:8.642 毫秒。
我们观察到从最初的 4 到 5 秒时间显著降低到仅 12 毫秒的索引。此外,搜索时间进一步降低到仅 8.6 毫秒的 HNSW 索引。
关于如何在本地快速部署 pgvector 或者在线体验 pgvector 威力,在上篇 的文末有介绍。
评论