翻译: ScaNN：高效向量相似性搜索

Posted 2022-02-17 AI架构师易筋

假设一个人想要使用需要完全匹配标题、作者或其他易于机器索引的标准的查询来搜索大型文学作品数据集。这样的任务非常适合使用 SQL 等语言的关系数据库。但是，如果想要支持更抽象的查询，例如“内战诗歌”，则不再可能依赖简单的相似性指标，例如两个短语之间的共同词数。例如，查询“科幻小说”与“未来”的相关性比与“地球科学”的相关性更高，尽管前者有 0，而后者有 1，与查询相同的词。

机器学习 (ML) 极大地提高了计算机理解语言语义并因此回答这些抽象查询的能力。现代 ML 模型可以将文本和图像等输入转换为嵌入、高维向量训练，从而使更多相似的输入更紧密地聚集在一起。因此，对于给定的查询，我们可以计算其嵌入，并找到其嵌入最接近查询的文学作品。通过这种方式，ML 将一项抽象且以前难以指定的任务转变为一项严格的数学任务。然而，一个计算挑战仍然存在：对于给定的查询嵌入，如何快速找到最近的数据集嵌入？嵌入集通常对于穷举搜索来说太大了，而且它的高维使得修剪变得困难。

在我们的ICML 2020论文“使用各向异性向量量化加速大规模推理”，我们通过关注如何压缩数据集向量以实现快速近似距离计算来解决这个问题，并提出了一种新的压缩技术，与以前的工作相比，该技术显着提高了准确性. 这项技术在我们最近开源的向量相似性搜索库(ScaNN) 中得到了应用，并且使我们能够在ann-benchmarks.com上测量出两倍于其他向量相似性搜索库的性能。

向量相似性搜索的重要性

基于嵌入的搜索是一种有效地回答依赖于语义理解而不是简单的可索引属性的查询的技术。在这种技术中，机器学习模型被训练以将查询和数据库项目映射到公共向量嵌入空间，使得嵌入之间的距离具有语义意义，即相似的项目更接近。

上图所示的双塔神经网络模型是一种特定类型的基于嵌入的搜索，其中查询和数据库项通过两个相应的神经网络映射到嵌入空间。在此示例中，模型响应对假设文学数据库的自然语言查询。

要使用这种方法回答查询，系统必须首先将查询映射到嵌入空间。然后它必须在所有数据库嵌入中找到最接近查询的嵌入；这是最近邻搜索问题。定义查询-数据库嵌入相似度的最常见方法之一是通过它们的内积；这种最近邻搜索称为最大内积搜索（MIPS）。

由于数据库大小很容易达到数百万甚至数十亿，MIPS 往往是推理速度的计算瓶颈，穷举搜索是不切实际的。这需要使用近似 MIPS 算法，这些算法可以用一些准确性来换取比蛮力搜索的显着加速。

一种新的 MIPS 量化方法

MIPS的几个最先进的解决方案基于压缩数据库项目，以便可以在蛮力所花费的时间的一小部分内计算出它们的内积的近似值。这种压缩通常通过学习量化来完成，其中从数据库中训练出矢量码本，并用于近似表示数据库元素。

以前的向量量化方案对数据库元素进行了量化，目的是最小化每个向量x与其量化形式x̃之间的平均距离。虽然这是一个有用的指标，但对此进行优化并不等同于优化最近邻搜索精度。我们论文背后的关键思想是更高的平均距离实际上可能会导致更高的 MIPS 精度。

基于这种直觉，我们更严厉地惩罚与原始向量平行的量化误差。由于其损失函数的方向依赖性，我们将我们的新型量化技术称为各向异性矢量量化。该技术能够以较低内积增加的量化误差来换取高内积的卓越精度，这是关键创新和性能提升的源泉。

ScaNN 中的

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

ScaNN 是开源软件，您可以在GitHub 上亲自试用。该库可以通过 Pip 直接安装，并具有 TensorFlow 和 Numpy 输入的接口。有关安装和配置 ScaNN 的更多说明，请参阅 GitHub 存储库。

结论

通过修改向量量化目标以符合 MIPS 的目标，我们在最近邻搜索基准上实现了最先进的性能，这是基于嵌入的搜索性能的关键指标。尽管各向异性矢量量化是一项重要的技术，但我们认为这只是通过优化算法以提高搜索精度的最终目标而不是压缩失真等中间目标可实现的性能增益的一个示例。

致谢

这篇文章反映了整个 ScaNN 团队的工作：David Simcha、Erik Lindgren、Felix Chern、Nathan Cordeiro、Ruiqi Guo、Sanjiv Kumar 和 Zonglin Li。我们还要感谢 Dan Holtmann-Rice、Dave Dopson 和 Felix Yu。

参考

https://ai.googleblog.com/2020/07/announcing-scann-efficient-vector.html