揭开 ScaNN 的神秘面纱：高效的向量相似性搜索

Posted 2021-06-29 TensorFlow 社区

 

文 / 软件工程师 Philip Sun，Google Research

 

人们要从大量的文学作品中进行搜索，一般需要使用与标题、作者或其他易于机器索引的标准完全匹配的查询条件。那么使用 SQL 等语言的关系型数据库可以轻松完成此类任务。

但是，如果想要查询如 ‘内战期间的诗歌’ 这类比较抽象的内容，就无法再依赖如对比短语之间共同单词数量等简单的 相似性指标。如，“科幻小说 (Science Fiction)”与“未来 (Future)”的关系比与“地球科学 (Earth Science)”的关系更紧密，尽管，后者有一个相同单词  Science，而前者没有。

机器学习 (ML) 极大地提高了计算机理解语言语义以及响应这些抽象查询的能力。现代 ML 模型可以将文本和图像等输入转换为嵌入向量，对高维向量进行训练，使更相似的输入更紧密地聚集在一起。因此，对于给定的查询，我们可以计算其嵌入向量，并找到嵌入向量与查询最接近的文学作品。通过这种方式，ML 已经将以前难以指定的抽象任务转变为严格的数学任务。

然而，计算难题依然存在：对于给定的查询嵌入向量，如何快速找到数据集中最接近的嵌入向量呢？嵌入向量集合通常太大，无法进行穷尽式搜索，并且通常嵌入向量的维数很高。这使得很难对向量本身进行修剪。

在我们的 ICML 2020 论文“Accelerating Large-Scale Inference with Anisotropic Vector Quantization”中，我们通过关注如何压缩数据集向量来实现快速近似距离计算解决这一问题，并提出了一种新的压缩技术，与以前的工作相比，这项技术可以大大提高准确率。我们在近期开源的向量相似性搜索库 (ScaNN) 中应用了这项技术，与其他向量相似性搜索库相比，我们的性能高出两倍（在 ann-benchmarks.com 上测得）。

 

向量相似性搜索的重要性

基于嵌入向量的搜索技术能够有效响应依赖于语义理解的查询，这比简单的可索引属性的查询更为高效。在这种技术中，机器学习模型被训练成将查询和数据库项映射到公共向量嵌入空间，使得嵌入向量之间的距离带有语义含义，即相似项更接近。

要用这种方法响应查询，系统必须先将查询映射到嵌入空间。然后是最近邻搜索 (Nearest Neighbor Search) 问题：它需要在所有数据库嵌入向量中找到最接近查询的嵌入向量。定义查询数据库嵌入向量相似性的最常见方法之一是通过它们的内积；这种最近邻搜索被称为最大内积搜索 (Maximum Inner-Product Search，MIPS)。

数据库的大小很容易达到数百万甚至数十亿量级，这种情况下 MIPS 通常是推理速度的计算瓶颈，进行穷尽式搜索不切实际。而使用近似的 MIPS 算法虽然会牺牲一些精度，但能大幅提高蛮力搜索速度。

 

MIPS 的新量化方法

MIPS 的几种最新解决方案均基于压缩数据库项，可以超越蛮力搜索，迅速计算出它们的内积近似值。这种压缩通常通过学习量化 (Learned Quantization) 来完成，其中向量的 码本 (Codebook) 基于数据库训练得出，并用于近似表示数据库元素。

先前的向量量化方案对数据库元素进行了量化，目的是最小化每个向量 x 及其量化形式 x̃ 之间的平均距离。尽管这是一个有用的指标，但为此进行优化并不等同于优化最近邻搜索精度。论文背后的关键理论是，编码的平均距离 越大 ，实际上 MIPS 精度也可能越高。

我们对结果的分析如下所示。假设我们有两个数据库嵌入向量 x1 和 x2，并且必须将每个量化到两个中心之一：c1 或 c2。我们的目标是将每个 xi 量化为 x̃i ，以使内积 <q, x̃i> 尽可能接近原始内积 <q, xi>。可将此过程可视化为使 x̃i 在 q 上的投影幅度尽可能类似于 x̃i 在 q上的投影幅度。在传统的量化方法（左）中，我们将为每个 xi 选择最近的中心，这将导致两个点的相对排名不正确：<q, x̃1> 大于 <q,  x̃2>，但是实际上 <q, x1> 小于 <q, x2>！如果我们改为将 x1 赋给 c1，将 x2 赋给 c2，则会得到正确的排名。如下图所示。

目标是将每个 xi 量化为 x̃i =  c1 或 x̃i = c2：传统的量化（左）导致该查询的 x1 和 x2 的排序不正确。尽管我们的方法（右）选择的中心距数据点较远，但这实际上降低了内积误差、提高了精度

 

事实证明，方向和幅度同样重要——即使 c1 与 x1 的距离超过 c2 与它的距离， c1 在几乎与 x1 完全正交的方向上偏离 x1，而 c2 的偏离是平行的（对于 x2，情况相同，但方向相反）。平行方向上的误差在 MIPS 问题中危害更大，因为它严重影响高内积，根据定义，高内积是 MIPS 试图精确估计的内积。

基于这种认知，我们会更加严厉地惩罚与原始向量平行的量化误差。由于其损失函数的方向依赖性，我们将这种新型量化技术称为 各向异性向量量化 (Anisotropic Vector Quantization) 。它能够用较低内积增加的量化误差来换取较高内积的超高精度，这是创新的关键所在，也是其性能提升的源泉。

在上图中，椭圆表示等损失的轮廓：在各向异性向量量化中，平行于原始数据点 x 的误差会受到更大的惩罚

 

ScaNN 中的各向异性向量量化

各向异性向量量化使 ScaNN 可以更好地估计前 k 个 MIPS 结果中的内积，从而实现更高的精度。在来自 ann-benchmarks.com 的 glove-100-angular 基准测试中，ScaNN 的性能优于其他 11 个经过精心调整的向量相似性搜索库，在给定的精度下，每秒处理的查询数量大约是第二快的库的两倍。*

Recall@k 是最近邻搜索精度的常用指标，用于衡量算法返回的 k 个近邻中存在的真正的最近 k 个近邻的比例：ScaNN（上面的紫色线）在速度-精度权衡的各个方面始终展现出优异的性能

 

我们已经将 ScaNN 在 GitHub 上开源。ScaNN 可以通过 Pip 直接安装，并具有用于 TensorFlow 和 Numpy 输入的接口。有关安装和配置 ScaNN 的更多说明，请参见 GitHub 代码库。

 

结论

通过修改向量量化目标以与 MIPS 的目标保持一致，我们在最近邻搜索基准上获得了 SOTA 性能，这是基于嵌入向量的搜索性能的关键指标。

各向异性向量量化是一项重要技术，我们今天主要讨论的是通过优化算法实现性能提升的一个例子，其最终目标是实现提高搜索精度，而不是压缩失真等中间目标。

 

致谢

这篇文章是整个 ScaNN 团队共同努力的成果：David Simcha、Erik Lindgren、Felix Chern、Nathan Cordeiro、Ruiqi Guo、Sanjiv Kumar 以及 Zonglin Li。我们还要感谢 Dan Holtmann-Rice、Dave Dopson 和 Felix Yu。

 

* ScaNN 在 ann-benchmarks.com 的其他数据集上也有类似的出色表现，但该网站目前显示的是过时、较低的数据。

有关其他数据集上更具代表性的性能数据，请参阅此 PR (https://github.com/erikbern/ann-benchmarks/pull/172)。

 

登陆官网了解更多 TensorFlow 搭建推荐系统内容，可登录 B 站 Google 中国 TensorFlow 频道观看《使用 TensorFlow 搭建推荐系统系列（全八讲）》视频，也可关注 TensorFlow 官方公众号获取更多资讯。