加速多图向量搜索

原创

点火三周

发布于 2024-03-25 18:35:39

8540

发布于 2024-03-25 18:35:39

文章被收录于专栏：Elastic Stack专栏

加速多图向量搜索

Lucene中多图向量搜索的先前状态

如我们之前所述， Lucene 以及 Elasticsearch 的近似 kNN 搜索基于在 HNSW 图中搜索每个索引段并组合所有段的结果来查找全局 k 个最近邻。当最初引入时，多图搜索是在单个线程中顺序执行的，一个接一个地搜索每个段。这带来了一些性能损失，因为搜索单个图的大小是亚线性的。在Elasticsearch 8.10中，我们并行化了向量搜索，如果线程池中有足够的可用线程，则在 kNN 向量搜索中为每个段分配一个线程。由于这一变化，我们在夜间基准测试中看到查询延迟下降到之前值的一半。

尽管我们在段上并行搜索，但它们仍然是独立搜索，每个搜索都收集自己的前k个结果，而不与其他段同步搜索进展。从我们对词汇搜索的经验来看，我们知道通过在段搜索之间交换到目前为止收集的最佳结果信息，我们可以实现显著的搜索加速，我们认为我们可以将相同的想法应用于向量搜索。

通过在段搜索之间共享信息来加速多图向量搜索

当我们使用基于图的系统（比如HNSW）来寻找一个点的最接近的邻居时，其实是在用两种策略：一种是广泛探索，另一种是针对性利用。以HNSW为例，它不仅仅满足于找到最接近的几个点（top-k），而是会先找到一个更大范围内相对较近的点集（top-n）。这样做的时候，它会查看许多可能的路径，即使有些路径看起来一开始并不那么有希望。这种做法的好处是，它可以帮助系统从误入的死胡同中走出来，找到更好的结果。但这也意味着它会花时间探索一些最终并不会被选中的选项。

另一种策略更为直接，它专注于尽快减少与目前已知的第k个最接近点的距离，只关注那些有望成为最终答案的路径。

简而言之，就像在一个巨大迷宫中寻找出口，HNSW既会尝试一些看起来不那么直接的路线，以防错过更好的路径，也会利用已知的信息尽快接近目标，这样的双重策略使它能在寻找最近邻居时更加高效和准确。

因此，扩展匹配集的大小（top-n）是一个超参数，通过增加或减少邻近图中的探索，允许人们用运行时间来交换召回率。

正如我们已经讨论的，Lucene为数据的不同分区构建了多个图。此外，在大规模时，如果想要在几台机器上水平扩展检索，必须对数据进行分区并构建单独的图。因此，一个普遍有趣的问题是“在同时搜索多个图的最近邻的情况下，应该如何适应这种策略？”

当我们同时搜索多个图，并从每个图中挑选出最靠前的k个结果时，我们发现召回率会显著提高。这其实很好理解：因为在搜索更多不同的图时，我们探索了更多的可能性，这减少了我们仅仅困在某个局部最优解中的风险。但这种方法的代价是，它会让搜索过程变长。我们理想中的情况是，无论数据怎么分片（或说是分成了多少个图），召回率都能保持高效，同时搜索速度也能更快。

影响多图搜索和单图搜索效率的有两大因素：单图中存在的特定连接，以及多个独立的前n个结果集合。通常，除非把数据严格分割开，每个小图里的点与它的邻居比起来，在整体大图中可能只占了一小部分真正的最近邻。这就意味着，当你在多个图中搜索时，你可能会花费时间去探索那些实际上并不会是最终结果的点。由于每个小图都是独立创建的，所以必然会有一些“结构性”的成本。但好消息是，正如我们即将展示的那样，通过在不同搜索之间智能共享信息，我们可以减少由此带来的成本。

有了一个大家共享的全局top-n结果集之后，一个很自然的问题就是，我们该如何处理那些不太可能成为最终结果的图的部分呢？特别是那些它们的末端顶点比全局目前已知的第n差的匹配还要差的边。如果我们只是在搜索单一图，这些边是不会被考虑进去的。但考虑到每次搜索都从不同的地方开始，进展速度也不一样，如果我们用同样的规则来处理多图搜索，可能会导致搜索过早地结束，错过一些实际上非常接近查询点的邻居。下面的插图就是这个情况的一个示例。

图1 两个图片段显示了收集top-2 集合的同时搜索的快照。在这种情况下，如果我们要修剪未访问的末端顶点不具有全局竞争力的边，我们将永远不会遍历红色虚线边，也无法找到图 2 中所有的最佳匹配。

为了解决这个问题，我们设计了一个简单的方法，能够根据每次局部搜索是否在全局范围内具有竞争力，有效地在不同的搜索参数之间切换。为了实现这一点，除了定期同步的全局队列之外，我们还为每个局部图维护了两个局部队列，记录了距离查询点最近的向量的距离。一个队列的大小为n，另一个的大小为⌊g×n⌋。这里的g控制了非竞争性搜索的贪婪程度，是一个小于1的数值。实际上，g是一个我们可以自由调整的参数，用于平衡召回率和搜索速度。

随着搜索的进行，我们在决定是否遍历一个边时检查两个条件：i)如果我们单独搜索图时，是否会遍历这个边，ii)这条边的端顶点是否全局具有竞争力或者它是否在局部与“贪婪”的最佳匹配集有竞争力。具体来说，如果我们用q表示查询向量，候选边的终点向量为Ve，第n个局部最佳匹配为Vn，第⌊g×n⌋个局部最佳匹配为Vg，以及第n个全局最佳匹配为Vgb，那么我们会把ve加入搜索集合:

d(ve, q) < d(vn, q) AND (d(ve, q) < d(vg, q) OR d(ve, q) < d(vgb, q))

这里，d(⋅,⋅)表示索引距离度量。注意这个策略确保我们总是继续搜索每个图到任何局部最小值，并且根据g的选择我们仍然逃离一些局部最小值。

忽略一些关于同步、初始化等的细节，这就描述了对搜索过程的修改。正如我们展示的，这个简单的方法显著改善了搜索的延迟，并且在召回率上，虽然接近，但仍优于单图搜索。

对性能的影响

我们的夜间基准测试显示，在索引时并发运行的向量搜索查询速度提高了高达60%（平均查询延迟从54毫秒降低到32毫秒）。

图2 在升级到Lucene 9.10后，与索引同时运行的查询延迟显著下降，该版本包含了新的改变。

在索引外运行的查询我们观察到适度的加速，主要是因为数据集不是很大，包含2百万96维向量跨越2个分片（图3）。但对于这些基准测试，我们可以看到图中访问的顶点数量显著减少，因此向量操作的数量也减少了（图4）。

图3 我们发现不使用并发索引运行的查询延迟略有下降，特别是检索前 100 个匹配项时，矢量操作的数量（图 4）大幅减少。

图4 我们看到检索top-10和top-100匹配时使用的向量操作数量大幅减少。

加速效果在包含更高维度向量的较大索引上应该更加明显：在测试中，我们通常看到了2倍到3倍的加速，这与我们上面看到的向量比较次数减少是一致的。例如，下面我们展示了在Lucene夜间基准测试中向量搜索操作的加速情况。这些测试使用了768维的向量。值得一提的是，在Lucene基准测试中，向量搜索在单个线程中顺序处理一个接一个的图，但这次改变也对这种情况产生了积极影响。这是因为，在第一个图搜索之后收集的全局top-n集合为后续的图搜索设定了阈值，如果这些图中不包含有竞争力的候选者，就允许它们更早结束。

图5 该图显示，随着2月7日的改变提交，每秒查询数量从104查询/秒增加到219查询/秒。

对召回率的影响

多图搜索加速以稍微降低的召回率为代价。这是因为我们可能会停止探索一个基于其他图的全局匹配可能仍有更好匹配的图。关于召回率下降，有两点说明：i) 根据我们的实验结果，召回率仍然高于单个图搜索的召回率，就好像所有段都合并成了一个单独的图（图6所示）。ii) 我们的新方法在相同召回率下实现了更好的性能（Pareto优势）：它在性能上优于我们之前的多图搜索策略（图7所示）。