SOAR:使用 ScaNN 实现更快矢量搜索的新算法

@TXENM}E$04%H]9}$)HD{SB.png

SOAR 是对向量搜索的算法改进,它为 Google 的向量搜索库 ScaNN 引入了有效且低开销的冗余,从而使 ScaNN 更加高效。

高效的向量相似性搜索对于 Google 的许多机器学习 (ML) 应用至关重要。它通常用于搜索嵌入,即现实世界实体(如图像、网站或其他媒体)的向量表示。这些向量表示允许计算机以数学方式比较并找到对象之间的相似性,例如将同一主题的肖像归为一组。一旦嵌入的数据集变得太大,无法使用将查询与数据集中的每个嵌入进行比较的强力方法,就需要更高效的向量相似性搜索方法来进一步扩展。

为了突出 Google 在向量搜索算法方面的创新,ScaNN 向量搜索库于 2020 年开源(更多信息请参阅之前的研究博客文章)。从那时起,更大的数据集和向量搜索在新用例(例如检索增强生成)中的普及都推动了对更具可扩展性的向量搜索算法的需求。

自发布以来,ScaNN 一直得到积极的维护和改进,但今天我们特别高兴地宣布 ScaNN 的一项重大算法进步:正交性放大残差溢出 (SOAR)。在NeurIPS 2023上发表的“ SOAR:近似最近邻搜索的改进索引”中,我们描述了如何在 ScaNN 的向量索引中引入一些冗余来提高向量搜索效率,同时尽量减少对索引大小和其他关键向量索引指标的影响。SOAR 帮助 ScaNN 满足对向量搜索库不断增长的扩展和性能要求。

上图最初出现在2020 年 ScaNN 开源发布博客文章中,展示了嵌入向量搜索如何帮助回答假设文学数据库上的自然语言查询。今天,我们展示的研究使 ScaNN 能够更快地执行向量搜索。

SOAR 和冗余

SOAR 的关键直觉是将一些数学设计和实施优化的冗余引入 ScaNN 的近似向量搜索例程中。冗余是使用多个副本的概念,当另一个副本发生故障时,每个副本都可以作为备份,以降低给定系统完全失效的可能性。冗余用于许多工程学科,从拥有可以调整机翼襟翼以确保飞机稳定性的多个系统,到跨多个物理驱动器存储数据副本以防止数据丢失。另一方面,当冗余的副本容易受到常见的相关故障的影响时,冗余可能会提供一种虚假的安全感——例如,将数据的多个副本存储在同一个数据中心仍然会使它们面临该地区常见的自然灾害威胁,因此不如地理上分散的数据副本那么可靠。

SOAR2-地图

冗余的一个常见示例是跨全球多个数据中心复制数据以降低数据丢失的可能性。SOAR 使用冗余来降低无法找到最近邻向量的风险。

在近似向量搜索中,“失败”意味着无法找到查询的真正最近邻居,而是检索不太相似的向量。根据上述主题,SOAR 的冗余性使 ScaNN 不太可能错过最近邻居,而 SOAR 的具体数学公式(如下所述)可最大限度地减少相关失败,否则会损害搜索效率。因此,SOAR 提高了以固定搜索成本可实现的搜索精度,或者等效地降低了实现相同搜索精度所需的搜索成本。

SOAR:数学细节

为了突出 SOAR 的不同之处,我们首先描述了 ScaNN 不使用 SOAR 的近似搜索范式。我们重点介绍 ScaNN 如何解决向量相似性搜索的最大内积搜索(MIPS) 变体,该变体将向量相似性定义为查询与数据库向量的内积。这是因为 SOAR 以 MIPS 为目标,也因为 MIPS 适用于各种其他相似性,包括余弦和欧几里得距离。

在 ScaNN 能够回答向量搜索查询之前,它首先要经过索引阶段,在此阶段,输入数据集经过预处理以构建有助于高效向量搜索的数据结构。此索引步骤的关键部分是通过k均值进行聚类;在 SOAR 之前,聚类的执行方式是将数据集中的每个向量分配到一个k均值聚类中。每个聚类都有一个聚类中心,它充当分配给该聚类的向量的粗略近似值;仅当聚类中心位于距离查询向量最近的N 个中心中时,才会对聚类中的向量进行评估(N是一个参数;N越高,搜索精度越高,但计算成本也越高)。

动画演示了如何在 SOAR 之前执行聚类以在 ScaNN 中执行有效的矢量搜索。

为了说明这种算法何时难以找到最近邻居,考虑将向量x分配给中心为c的聚类。将x和c之间的差异表示为残差或r。对于查询q,查询向量相似度 ⟨ q , x ⟩ 与估计相似度 ⟨ q , c ⟩ 之间的差异为 ⟨ q , x ⟩ - ⟨ q , c ⟩ = ⟨ q , x - c ⟩ = ⟨ q , r ⟩,当r与q平行时,该差异最大,如下所示:

SOAR4-低估Tri

当查询 q 与 r 平行时,估计的内积 ⟨q, c⟩ 大大低估了真实的内积 ⟨q, x⟩。在这些情况下,SOAR 有助于找到最近邻居,否则往往会降低向量搜索的准确性。

在这种情况下,如果x是q的最近邻居,则x通常很难找到,因为尽管 ⟨ q , x ⟩ 很高,但 ⟨ q , r ⟩ 也很高,导致查询中心相似度 ⟨ q , c ⟩ 很低,所以这个特定的聚类很可能会被修剪而不再进行进一步搜索。 有很多方法可以缓解这个问题。 例如,计算更高质量的聚类往往会降低 r 的幅度,从而降低平均估计误差,而各向异性矢量量化(AVQ)会“塑造”误差,使得当查询与x不相似时误差往往最大,因此不太可能影响结果。

SOAR 通过采用完全不同的方法解决了这个问题:允许将向量分配给多个集群。直观地看,这在冗余原则下是有效的:当主要分配表现不佳时(当q与主要分配的r高度平行时),次要分配可以充当“备用集群”,以促进高效、准确的向量搜索。

这种冗余源于第二个分配提供了一个新的向量中心差r '。只要当 r 与q接近平行时这个r ' 不与q接近平行,这个次级中心就应该有助于 ScaNN 找到与q最近的邻居。然而,SOAR 比这种简单的冗余更进一步,修改了二次分配的分配损失函数,以明确优化独立有效的冗余:它旨在找到r ' 垂直于r的二次聚类,这样当q与r接近平行且主中心误差较大时,q将与r '接近正交,且次级中心误差较小。SOAR 修改后的损失效果如下所示:

SOAR 不仅通过二次分配引入冗余,而且还通过修改损失来最大化冗余的有效性。未修改的损失(首先显示)通常会导致无效的二次分配,例如 c' 具有较高的误差。SOAR 的修改损失(稍后)选择了 c'',这是一个更好的选择。

SOAR 的实施和理论分析还有更多细节,但 SOAR 背后的核心思想如上所述:使用修改后的损失将向量分配给多个集群。该技术之所以命名为 SOAR,是因为先前的研究将多重分配描述为“溢出”,而修改后的损失鼓励正交(垂直)残差,因此得名“正交性放大残差的溢出”。

实验结果

SOAR 使 ScaNN 能够保持其现有的优势,包括低内存消耗、快速索引速度和硬件友好的内存访问模式,同时赋予 ScaNN 额外的算法优势。因此,ScaNN 在向量搜索性能的三个主要指标之间取得了最佳平衡,如下面ann-benchmarks glove-100 数据集所示。唯一接近 ScaNN 查询速度的库需要超过 10 倍的内存和 50 倍的索引时间。同时,ScaNN 实现的查询吞吐量比具有同等索引时间的库高出几倍。

SOAR6-结果英雄

在所有经过基准测试的库中,带有 SOAR 的 ScaNN 实现了迄今为止最佳的查询速度/索引速度平衡,同时占用的内存最小。

ScaNN 还在其适用的Big-ANN 2023 基准测试的两个轨道中取得了最先进的性能。这些基准测试涉及更大的数据集,这实际上进一步提高了 SOAR 的效率。本文进一步讨论了数据集大小对 SOAR 的影响,以及对多达十亿个向量数据集的基准测试。

SOAR7-结果栏

对于 Big-ANN 2023 的分销外和流媒体赛道,SOAR 可让 ScaNN 在每个赛道的各自排名指标中取得最高成绩。该网站的排行榜正在等待更新。

结论

当 ScaNN 传统的基于聚类的方法最困难时,SOAR 为 ScaNN 提供了一条强大的“备用”路线来识别最近邻居。这使 ScaNN 能够执行更快的向量搜索,同时保持较低的索引大小和索引时间,从而在向量搜索算法中实现最佳的权衡。

我们邀请社区使用 ScaNN 来解决其向量搜索挑战。ScaNN 在GitHub上开源,可以通过 Pip 轻松安装。此外,ScaNN 向量搜索技术在 Google Cloud 产品中也可用:Vertex AI Vector Search利用 ScaNN 提供完全托管、高规模、低延迟的向量相似性匹配服务,而AlloyDB最近推出了ScaNN for AlloyDB index — 一个基于流行的 PostgreSQL 兼容数据库的向量数据库。我们很高兴看到更高效的向量搜索如何支持下一代机器学习应用程序。

致谢

这篇文章体现了整个 ScaNN 团队的辛勤劳动:David Simcha、Felix Chern、Philip Sun、Ruiqi Guo、Sanjiv Kumar 和 Zonglin Li。我们还要感谢 Apurv Suman、Dave Dopson、John Guilyard、Rasto Lenhardt、Susie Flynn 和 Yannis Papakonstantinou。

版权声明

本文仅代表作者观点,不代表本站立场。
本文系作者授权发表,未经许可,不得转载。

评论