Faiss向量召回引擎如何做到快速查找最近邻

欢迎来我的小站转转~~

背景

Faiss是Facebook开源的向量召回引擎，用于寻找与某个向量最相似的N个向量。

Faiss第一次release发布于2018.02.23，但其作者Matthijs在加入Facebook之前的2011年就已经发表了一篇关于最近邻搜索的论文，Faiss就是基于此论文思想实现的。读懂了这篇论文，Faiss的索引方式就清楚了。

问题描述

给定D维向量$x$和集合$\Gamma =y_1,y_2\dots y_N$ ,需要找到与$x$距离最短的k个最近邻。

以欧氏距离为例，可表示为：

$$L=k-argmi{n}_{i=0:N}||x-y_i||$$

在我们的应用中，$x\in \Gamma$

问题规模

我们试着以最粗暴的方法进行穷举搜索，来看一下这个解的复杂度有多高。

1. 构造距离矩阵：每两个向量$x$与$y$的距离计算公式为$\sqrt{\Sigma (x_i-y_i{)}^2},i\in [1,D]$ ，耗费时间为 $O(D)$ , 距离矩阵包含 $N^2$ 个元素，总共耗时为 $O(D\ast N^2)$
2. 从距离矩阵中查找到k个最近邻，若用最小堆算法，时间复杂度为$O((N-k)logk)$

取$N=2000W,k=1000,D=1000$

得到

• 距离矩阵包含400T个元素，假设每个距离为float占用32bit，至少占用1600TB空间
• 构造距离矩阵运算时间复杂度数量级为$10^{17}$
• 从距离矩阵中找到k个最近邻的时间复杂度数量级为$10^9$

最近邻离线表

一般来说向量集合都是每天更新的，这时候可以试着直接把每个向量对应的k个最近邻保存起来

1. 构造距离矩阵+N个k近邻查找耗时为 $O(D{N}^2+N\ast (N-k)\ast logk)$
2. 构造最近邻离线表空间占用 $N\ast k$

在我们的场景中得到

1. 找到N个k近邻查找的耗时数量级为 $10^{17}$
2. 存储空间为 $1600TB$ (构造后删除)+320G(假设每个索引用int表示，分数用float表示)

查找k最近邻的耗时: O(1)

向量量化(Vector Quantization)

在我们的场景中，其实不需要最精确的距离，允许一定程度的误差。在这种情况下，我们可以引入向量量化方法，将向量的数量大幅度缩小。

所谓向量量化，就是将原来无限的空间 $R^D$ 映射到一个有限的向量集合 C = { c i , i ∈ [ 1 , l ] } \mathcal{C} = \{c_i, i\in[1,l]\} C={ ci​,i∈[1,l]} 中，其中 $l$ 是一个自然数。将这个从 $R^D$ 到集合 $\mathcal{C}$ 的函数记为 $q$ ，则 $\forall q(y)\in \mathcal{C}$，在信息论中称 $\mathcal{C}$ 为codebook。

当然这里的映射函数也不是随便指定的，需要满足误差最小的原则，一种方法是将优化函数设置为最小平方误差$MSE(q)={\mathbb{E}}_X[d(q(y),y{)}^2]$

咦，正好就是k-means方法的目标函数！因此我们可以用k-means作为寻找最佳codebook的方法。

那现在我们来分析一下进行向量量化占用的空间和时间复杂度

假设我们将原来2000W个向量映射到大小为20W的集合中(平均每个中心点代表100个向量，已经引入了较大的误差)

1. 距离矩阵：只需要存储对应$\mathcal{C}$ 中向量之间的距离，占用的空间为$||\mathcal{C}|{|}^2$，此例中为$400G\times 4B=1.6T$
2. $y\to c\in \mathcal{C}$的映射关系，若以int标识一个向量，则共约$N\times 4=80M$内存
3. 时间复杂度：k-means算法的时间复杂度为 O ( m i t e r N k D