Facebook 开源了 AI 相似性搜索工具 Faiss。而在一个月之后的今天,Facebook 发布了对 Faiss 的官方原理介绍。
它是一个能使开发者快速搜索相似多媒体文件的算法库。而该领域一直是传统的搜索引擎的短板。借助Faiss,Facebook 在十亿级数据集上创建的最邻近搜索(nearest neighbor search),比此前的最前沿技术快 8.5 倍,并创造出迄今为止学术圈所见最快的、运行于 GPU 的 k-selection 算法。Facebook 人工智能实验室(FAIR) 借此创造了数个世界纪录,包括在十亿高维矢量上的构建的、世界最快的 k-nearest-neighbor 图。
相似性搜索的本质
传统数据库由包含符号信息的结构表组成。比方说,一个图像集,会用每行放一张索引照片的列表来表示。每一行都包含诸如图像标识和描述语句等信息。每一行也可与其他表格的条目关联,比如照片与人名列表相关联。
雷锋网获知,很多AI 工具都会产生高维矢量,比如像 word2vec 这样的文本嵌入工具,以及用深度学习训练的 CNN 描述符(descriptors)。这些表示比固定的符号表示更加强大灵活,至于原因,本文将为大家解释。但是,用 SQL 来检索的传统数据库并没有适配这些新型表示。首先,海量的新多媒体流创造了数十亿的矢量。其次,而且更重要的是,找到相似的相似的条目意味着找到相近的高维矢量。而对于当下的标准检索语言,这是极度低效、甚至无法实现的。
如何使用矢量表示?
让我们假设你有一张某建筑的影像,比方说某城市的礼堂照片,但你忘记这是哪一个城市的了。然后,你希望找到图片库中该建筑的所有照片。该情况下,SQL 中常用的 key/value 检索并没有帮助——因为你已经忘了这是哪个城市。
这就轮到相似性搜索派上用场。由于设计,图像的矢量表示会对相似图像生成相近的矢量。在这里,相近矢量被定义为在欧几里得空间相邻的矢量。
另一项矢量表示的应用是分类。想象下你需要一个分类器,来判别图片库中哪一个图片代表了菊花。分类器的训练是一个开发者们都比较熟悉的过程:算法把菊花和非菊花的图像作为输入。若果分类器是线性的,它会输出一个分类矢量,后者带有一项重要属性:它的向量点积(Dot Product)和图像矢量在一起,能反映出该图像包含菊花的概率。然后对向量点积和图片库中的所有条目进行计算。最后 return 有最高概率值的图像。这种检索是一种“最大内积”搜索。
所以,对于相似性搜索和分类,我们需要以下操作:
给定检索矢量,return 在欧几里得距离上最接近这个矢量的数据库对象列表。
给定检索矢量,return 有最高向量点积的数据库对象列表。
一个额外的挑战,是 Facebook 想要在一个大尺度上执行这些操作。这里,“大尺度”指的是数十亿的矢量。
软件包
现有的软件工具,不足以完成上述数据库搜索操作。传统的 SQL 数据库系统可用性不高,因为它们是为 hash-based searches 或 1D interval searches 而优化。OpenCV 等工具包里包含的相似性搜索功能,在扩展性上的限制非常大。针对“小”数据集的相似性搜索算法库也是这么个情况(比如,一百万个矢量)。其他工具包大多是学术研究的产物,为发表的论文而开发,用来在特定设定下展示效果。
Faiss 是一个打破上述限制的算法库。它的优点有:
提供数个相似性搜索方法。这些方法针对不同使用情况,提供了跨度很大的功能取舍。
为内存的使用和速度而优化。
为相关索引方法提供了最前沿的 GPU 执行方案。
相似性搜索评估
一旦这些矢量被学习机提取出来(从图像、视频、文本文件或其他渠道),它们就已经可以被输入进相似性搜索库。
Facebook 有一个作为参照的“暴力”算法——它会完整地照章计算所有的相似性,然后 return 最相似元素的列表。这提供了一个黄金结果参照列表。值得注意的是,高效得执行该暴力算法不容易实现,而且经常影响系统其它部分的效果。
如果我们愿意牺牲一部分精确度,相似性搜索的速度可以有数个数量级的提升。当然,这会导致相对参照结果的一点偏移。举个例子,对图像相似性搜索的第一和第二个结果进行交换,或许不会有什么区别,因为它们很可能都是某个给定检索的正确答案。加速搜索意味着要对数据集进行一些预处理,Facebook 把这成为索引。
这使 Facebook 确定了三大研究方向:
速度。找到十个最相似的矢量需要多久?希望花费的时间比暴力算法要少。不然的话,索引就没有任何意义。
内存使用。该方法需要多少 RAM?比原始矢量多还是少? Faiss 只支持在 RAM 上搜索,因为其他磁盘数据库的速度要慢数个数量级。即便是 SSD 也太慢。
精确度。返回的结果列表与暴力搜索结果差多少?精确度能通过计算检索数量,在结果列表中先返回最邻近单位评估;或是衡量 10 个最先返回的最邻近单位的平均 fraction (该方法被称之为 10-intersection)。
Facebook 一般会衡量在给定内存使用情况下,速度和精确度之间的权衡。Faiss 专注于压缩原始矢量的方法,因为它们是扩展到十亿级矢量数据集的唯一途径。当需要索引的矢量有十亿个之多,每一个会占用 32 左右的字节,这些矢量会占用极大的内存空间。
许多索引算法库针对的是百万左右的矢量,Facebook 的工程师们把这成为小规模。举个例子,nmslib 就包含这方面极其高效的算法,它比 Faiss 更快但是会占用远远更多的内存空间。
十亿个矢量的评估
工程世界中,并没有针对这个大小数据集的基准。因此,Facebook 基于研究结果进行评估。
精确度在 Deep1B 上进行评估,它是包含十亿张图片的图像库。每一个图像都已经被 CNN 处理过,CNN 中的其中一个 activation map,会被作为图像 描述符(descriptor)。这些矢量可以与欧几里得距离进行比较,来量化这些图像之间的相似度。
Deep1B 包含一个比较小的检索图像库。真实的相似性搜索结果,由处理了这些图像的暴力算法提供。因此,如果我们运行一个搜索算法,我们就可以评估结果中的 1-recall@1。
选择索引
由于评估,我们把内存使用限制在 30 GB。该内存限制指导我们进行索引方法和参数的选择。在 FAISS,索引方法用字符串来表示;在这个例子中是OPQ20_80,IMI2x14,PQ20。
该字符串代表了应用于矢量的预处理步骤 (OPQ20_80) 。一个 selection mechanism (IMI2x14) 来指示数据库应该如何被分割,以及一个编码部分(encoding component PQ20)来指示用 product quantizer (PQ) 编码的矢量,生成 20 字节的代码。因此,内存的占用(包括了 overheads)在 30GB 以下。
这听起来太技术流,因此 Faiss 的文件会向开发者提供指导:如何根据需要选择最恰当的索引类型。
索引类型确定之后,就可以开始索引。FAISS 算法库对这十亿个矢量进行处理,并把他们放入索引。索引能够存在硬盘,或者立即使用,对索引的搜索、additions/removals 可被交错插入。
在索引中搜索
当索引就绪后,一系列 search-time 的参数可设为针对此方法进行调整。由于评估需要,我们用单线程进行搜索。由于内存占用已经被限制住,我们需要在精确度和搜索时间之间进行权衡、优化。举个例子,这意味着能对 1-recall@1 40% 的最不可能搜索时间设置参数。
幸运的是,Faiss 配有自动调参机制,能扫描参数空间,并对提供了最佳操作点(operating points)的那些进行。这意味着给定精确度情况下的最优潜在搜索时间,或者反过来,给定搜索时间的最优精确度。在 Deep1B 上,操作点可用折线图的形式进行可视化。
这幅图上,我们可读出,获取 40% 的 1-recall@1,有少于每矢量 2 ms的检索时间。如果把检索时间放宽到 0.5 ms,我们可以达到 30%。2 ms 的检索时间,意味着单核的 500 QPS(queries per second )。
若把该结果与圈内最先进的研究相比,即 Babenko 和 Lempitsky 在 CVPR 2016 上发表的论文:“Efficient Indexing of Billion-Scale Datasets of Deep Descriptors”。该论文介绍了 Deep1B 数据集。但他们需要 20 ms 来获取 45% 的 1-recall@1。
用 GPU 处理十亿级数据集
当前,许多研究努力集中于 GPU 的执行上。在原生多 GPU 支持下,这能够产生相当不错的单机性能。GPU 执行可被看做是对应 CPU 的替代,你其实不需要理解 CUDA API 来挖掘 GPU 的性能。 Faiss 支持所有 2012 年之后发布的英伟达显卡(开普勒,计算能力 3.5+)。
我们希望把 roofline model 作为指南,它指出开发者应尽量使内存带宽或者浮点单位饱和。Faiss 单 GPU 的速度一般比 CPU 快五到十倍。新的帕斯卡架构 GPU 硬件,比如英伟达 P100,使之快了 20 倍有余。
一些展示性能的数字:
合适的索引,一个简单暴力的 k-nearest-neighbor 图(k = 10),基于 YFCC100M 数据集中 9500 万图像的128D CNN 描述符,0.8 的 10-intersection,用四路上代泰坦(Maxwell Titan X)只需要 35 分钟即可建成,包含索引创建时间。
十亿矢量的 k-nearest-neighbor 图已经即将成为现实。开发者可以在 Deep1B 数据集上创建强力的 k-nearest-neighbor 图(k = 10),0.65 的 10-intersection,在四路 Maxwell 泰坦支援下需要 12 个小时。若是 0.8 的 10-intersection,八路帕斯卡 P100-PCIe GPU,也是 12 个小时。更低质量的图可在五小时内用泰坦创建完成。
其他部分也达到了惊人性能。比方说,创建上述 Deep1B 索引需要 6710 万 120-dim 矢量,用 k-means 聚类生成 262,144 个几何中心。这在 25 E-M 次迭代下,需要四路泰坦 (12.6 tflop/s of compute) 花 139 分钟进行处理。注意聚类的训练集不需要与 GPU 显存匹配,因为数据是按需即时导入 GPU 的,而不会影响性能。
底层技术
雷锋网获知,FAIR(Facebook 人工智能实验室) 自 2015 年着手开发 Faiss。这是建立在过去的许多研究结论和大量技术攻关的基础上。对于 Faiss,Facebook 选择专注于对几项基础技术进行优化。尤其在 CPU 方面,Facebook 大量利用了:
多线程以充分利用多核性能并在多路 GPU 上进行并行搜索。
BLAS 算法库通过 matrix/matrix 乘法进行高效、精确的距离计算。没有 BLAS,高效的强力执行很难达到最优状态。 BLAS/LAPACK 是唯一一个 Faiss 必须的前提软件。
机器 SIMD 矢量化和 popcount 被用于加速孤立矢量的距离计算。
在 GPU 方面
对于从前的相似性搜索 GPU 执行,k-selection(寻找 k-minimum 或 maximum 因子)一直存在性能问题。这是因为普通的 CPU 算法(比如 heap selection)并不适用于 GPU。对于 Faiss GPU,Facebook 设计了学术圈迄今为止最快的小型 k-selection 算法(k <= 1024)。所有中间状态都完全保存在寄存器中,进一步提升了速度。它能够将输入数据以 single pass 进行 k-select,运行于潜在峰值性能的 55%,取决于峰值 GPU 显存带宽。由于其状态只存储在注册表中,并可与其他 kernels 融合,使它成为超级快的 exact 和 approximate 搜索引擎。
研究领域的许多注意力被放到了高效的 tiling 策略,和面向 approximate 搜索的 kernels 执行。多 GPU 支持用过粉碎或复制数据来提供。开发者并不会受单 GPU 显存大小的限制。半精度浮点支持 (float16) 也有提供,可在支持的 GPU 上进行完整 float16 运算,或者更早 GPU 架构所提供的的中级 float16 存储。我们发现 float16 这样的编码矢量能在几乎不损失精度的前提下进行加速。
简而言之,持续的 overhead 因素会在执行中起到作用。Faiss 做了许多关注工程细节的痛苦工作。
上手 Faiss
Faiss 用 C++ 实现,支持 Python。想要上手的各位,请到 GitHub 获取 Faiss,进行编译,然后把 Faiss 模块导入到 Python。Faiss 与 numpy 能做到完美的整合,包括需借助 numpy 阵列来实现的所有功能 (in float32)。
GitHub 地址:https://github.com/facebookresearch/faiss
详情请访问:https://code.facebook.com/posts/1373769912645926/faiss-a-library-for-efficient-similarity-search/