宣布 ScaNN:高效的向量相似性搜索
假设有人想要使用需要精确匹配标题、作者或其他易于机器索引的条件的查询来搜索大型文学作品数据集。这样的任务非常适合使用 SQL 等语言的关系数据库。但是,如果想要支持更抽象的查询,例如“内战诗”,则不再可能依赖朴素的相似度指标,例如两个短语之间的共同词数。例如,查询“科幻”与“未来”的相关性比与“地球科学”的相关性更高,尽管前者与查询相同的词为 0,而后者只有一个。
机器学习 (ML) 极大地提高了计算机理解语言语义并回答这些抽象查询的能力。现代 ML 模型可以将文本和图像等输入转换为嵌入、高维向量训练,以便更多相似的输入聚集在一起。因此,对于给定的查询,我们可以计算其嵌入,并找到嵌入最接近查询的文学作品。通过这种方式,ML 将一项抽象且以前难以指定的任务转变为一项严格的数学任务。然而,一个计算挑战仍然存在:对于给定的查询嵌入,如何快速找到最近的数据集嵌入?对于穷举搜索而言,嵌入集通常太大,并且其高维度使得修剪变得困难。
在我们的ICML 2020论文“使用各向异性向量量化加速大规模推理”,我们通过关注如何压缩数据集向量以实现快速近似距离计算来解决这个问题,并提出了一种新的压缩技术,与之前的工作相比,该技术显着提高了准确性. 这种技术在我们最近开源的向量相似性搜索库(ScaNN) 中得到了利用,使我们能够在ann-benchmarks.com 上测得的性能比其他向量相似性搜索库高两倍。
矢量相似性搜索的重要性
基于嵌入的搜索是一种有效地回答依赖语义理解而不是简单的可索引属性的查询的技术。在这种技术中,机器学习模型被训练以将查询和数据库项目映射到公共向量嵌入空间,从而嵌入之间的距离带有语义意义,即相似的项目更接近。
要使用这种方法回答查询,系统必须首先将查询映射到嵌入空间。然后它必须在所有数据库嵌入中找到最接近查询的嵌入;这是最近邻搜索问题。定义查询-数据库嵌入相似度的最常见方法之一是通过它们的内积;这种最近邻搜索称为最大内积搜索(MIPS)。
由于数据库规模很容易达到数百万甚至数十亿,MIPS往往是推理速度的计算瓶颈,穷举搜索是不切实际的。这需要使用近似的 MIPS 算法,该算法将一些准确性交换为比蛮力搜索显着的加速。
一种新的 MIPS 量化方法 MIPS 的
几种最先进的解决方案基于压缩数据库项,以便可以在很短的时间内计算出它们的内积的近似值。这种压缩通常是通过学习量化完成的,其中向量的码本是从数据库中训练出来的,用于近似表示数据库元素。
以前的矢量量化方案量化了数据库元素,目的是最小化每个矢量x与其量化形式x̃之间的平均距离. 虽然这是一个有用的指标,但对此进行优化并不等同于优化最近邻搜索精度。我们论文背后的关键思想是,具有更高平均距离的编码实际上可能会导致更高的 MIPS 精度。
我们的结果的直觉如下所示。假设我们有两个数据库嵌入x 1和x 2,并且必须将每个嵌入到两个中心之一:c 1或c 2。我们的目标是将每个x i量化为x̃ i使得内积 < q , x̃ i> 尽可能类似于原始内积 < q , x i > 。这可以被可视化为使得突出部的大小X我到q尽可能相似的投影X我到q。在传统的量化方法(左)中,我们会为每个x i选择最近的中心,这会导致两个点的相对排名不正确:< q , x̃ 1 >大于< q , x̃ 2 >,甚至虽然 < q , x1 >小于< q , x 2 >!如果我们改为将x 1分配给c 1并将x 2分配给c 2,我们会得到正确的排名。这如下图所示。
事实证明,方向事项以及大小-即使Ç 1是从更远的X 1比Ç 2,c ^ 1从偏移X 1的方向上几乎完全正交于X 1,而c ^ 2的偏移量是平行(对于x 2,同样的情况适用但翻转)。平行方向的误差在 MIPS 问题中危害更大,因为它不成比例地影响高内积,根据定义,这是 MIPS 试图准确估计的内积。
基于这种直觉,我们更严厉地惩罚与原始向量平行的量化误差。由于其损失函数的方向依赖性,我们将我们的新型量化技术称为各向异性矢量量化。这种技术能够用较低内积的增加的量化误差来换取高内积的卓越精度,这是关键创新和其性能提升的来源。
ScaNN 中的
各向异性矢量量化 各向异性矢量量化允许 ScaNN 更好地估计可能出现在 top- k MIPS 结果中的内积,从而实现更高的精度。在来自ann-benchmarks.com的glove-100-angular 基准测试中,ScaNN 的表现优于其他 11 个经过精心调整的向量相似性搜索库,在给定精度下每秒处理的查询数量大约是次快库的两倍。*
[email protected] 是最近邻搜索准确度的常用指标,它衡量算法返回的 k 个邻居中存在的真实最近 k 个邻居的比例。ScaNN(上紫色线)在速度-精度权衡的各个点上始终如一地实现了卓越的性能。ScaNN 是开源软件,您可以在GitHub 上自行试用。该库可以通过 Pip 直接安装,并具有用于 TensorFlow 和 Numpy 输入的接口。有关安装和配置 ScaNN 的进一步说明,请参阅 GitHub 存储库。
结论
通过修改矢量量化目标以符合 MIPS 的目标,我们在最近邻搜索基准上实现了最先进的性能,这是基于嵌入的搜索性能的关键指标。尽管各向异性矢量量化是一项重要技术,但我们认为这只是通过优化算法以提高搜索精度的最终目标而不是诸如压缩失真等中间目标来实现性能提升的一个例子。