Faiss向量召回引擎如何做到快速查找最近邻
@ Shen Jianan · Sunday, Apr 26, 2020 · 9 minute read · Update at Apr 26, 2020

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

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

问题描述

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

以欧氏距离为例,可表示为:

$$L = k-argmin_{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*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(DN^2 + N*(N-k)*logk)$
  2. 构造最近邻离线表空间占用 $N*k$

在我们的场景中得到

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

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

向量量化(Vector Quantization)

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

所谓向量量化,就是将原来无限的空间 $R^D$ 映射到一个有限的向量集合 $\mathcal{C} = {c_i, i\in[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_{iter}NkD)$,其中$m_{iter}$为迭代次数,N为原空间向量数量,k为中心点数量,D为向量维度。在此例中时间复杂度为$4m*10^{15}$,取迭代次数为25就已经达到了暴力搜索的数量级,只要迭代次数稍微上升些,时间复杂度还会更高。

乘积量化PQ(Product Quantization)

很多时候我们向量不同部分之间的分布是不同的

那么就可以将向量分成$m$个不同的部分,对每个部分进行向量量化,假设平均划分,则每个部分的维度大小为$D^*=D/m$

一个向量 $[x_1,x_2,…,x_{D^*},…,x_{D-D^*+1},…,x_D]$ ,可以划分为m组向量$[x_{i_1},x_{i_2},…,x_{i_D^*}]$,每组的codebook为$\mathcal{C_i}$,对应的量化器记为$q_i$,$\forall q_i(x_{i*}) \in \mathcal{C_i}$。则最终的全局codebook就是 $\mathcal{C} = \mathcal{C_1} * \mathcal{C_2} …* \mathcal{C_m}$,乘积量化的名称也来源于此。

以m=4为例,要达到上一节的20W量级,$|| \mathcal{C_i} ||$只需要达到22即可,要恢复到2000W级别也只需要达到67即可。

以 $\mathcal{C}$ 达到2000W量级的情况为例,每个分组的codebook中心点数量 $k^*$ 为67,现在来分析下对应的空间和时间复杂度

  1. 中心点之间的距离矩阵: 记 $||\mathcal{C_i}||$ 为 $k^{*}$ ,距离矩阵大小为 $O(m*(k^{*})^2)$ 。若用float表示距离,此例子中距离矩阵约为70M
  2. $y \to c \in \mathcal{C_i}$的映射关系,若以int标识一个向量,则共约 $m \times N \times 4=320M$内存
  3. 聚类时间复杂度:k-means算法的时间复杂度依然为$O(m_{iter} NkD)$,在此例中需要对m组分别进行k-means,总时间复杂度为$O(mNk^*D*m_{iter})$,此例中为$5m*10^{12}$,比上例降低了3个数量级
  4. 距离矩阵时间复杂度:维度D与距离矩阵元素数量的乘积,即 $O(Dm(k^*)^2)$,约$7 \times 10^{10}$

乘积量化大幅度降低空间占用的本质原因就在于:表达的向量空间是 $(k^*)^m$,但占用的磁盘空间为$mk^*$

论文中给出的经验取值是 $k^*=256$ ,$m=8$ ,对应的向量空间大小为$2^{64}$约$1.8\times 10^{19}$

Quantization场景下的距离计算

在没有Quantization的场景下,距离计算是直接对两个点计算$|| x-y || _2$

但在Quantization的场景下,不需要直接计算x和y的距离,而是通过中心点进行计算。这种方式有两个变种:

SDC

SDC(Symmetric Distance Computation): 两个向量之间的距离以两个向量所在的中心点距离来度量。误差小于等于x到中心点的距离+y到中心点的距离。

在PQ场景下,SDC距离表示为:$SDC (x,y)= \hat{d}(x,y)=d(q(x),q(y))=\sqrt{\Sigma_j{d(q_j(x),q_j(y))^2}}$

由于中心向量之间的距离已经存好(见上一节),计算x与y之间的距离只要查表即可,表的规模约70M,计算距离矩阵耗时数量级$10^{10}$

如果不想占用距离矩阵的空间,则时间复杂度为$O(Dmk^*)$,约为$10^5$

ADC

ADC(Asymmetric Distance Computation):x与y之间的距离以x与y所在的中心点距离来度量。用到三角形性质:两边之差小于第三边,所以误差一定小于等于y与中心点之间的距离

在PQ场景下,ADC距离表示为:$\tilde{d}(x,y)=d(x,q(y))$

记向量x的第i组元素为$u_i(x)$,则$ADC(x,y)=\sqrt{\Sigma_j{d(u_j(x),q_j(y))^2}}$

此场景下如果想要通过查表得到距离,则数据准备环节的距离矩阵的大小约为 $O(mnk^* )*4=20G$ (每个向量都要和m个分组中每个中心向量计算距离),耗时数量级 $10^{13}$

若直接计算,则复杂度与不查表版本的SDC相同,也是$O(Dmk^*)$,约为$10^5$

若使用查表策略,则ADC在精度更高的情况下,付出了更多的内存代价

IVFADC

上一节中,直接计算的时间主要耗费在与所有的中心向量进行对比上了。一种很自然的方法就是先找到一个大概的候选中心节点,避免与大量根本不可能的是最近邻的点进行计算。

粗糙量化(coarse quantization)+残差量化

因此,Matthijs在论文中提出了粗糙量化+残差量化的过程。具体来说,就是先从整个数据集合中构造一个大小为$kk^{'}$ (假设取值为1000)的小规模codebook $\mathcal{C_c}$,量化器记为$q_c$,于是每个向量都会有一个残差$r(y)=y-q_c(y)$。原始的向量可能会有特别大的分布差异/不平衡,但通过残差化之后的结果可以大幅度缓解这种问题。

再对$r(y)$使用PQ步骤,由于 $r(y)$ 相对于原始向量的"能量"更低,所以通过PQ步骤可以更精确地进行模拟。记PQ步骤的量化器为 $q_p$,则 y 通过 $q_c(y)+q_p(y-q_c(y))$ 来表示。这样的话两个向量x、y之间的距离 $d(x,y)$ 可以近似表示为

$$\ddot{d}(x,y) =d (x,q_c(y) + q_p(y-q_c(y))) = d(x - q_c(y) , q_p(y-q_c(y))) = \sqrt{\Sigma_jd({u_j(x-q_c(y)),q_{p_j}(u_j(y-q_c(y))))^2}}$$

记残差量化的codebook大小为$k_p$ (以64为例),如果要将这里的 $u_j(x-q_c(y))$ 提前计算好,即对每个x提前计算好与所有中心向量 $c \in \mathcal{C_c}$ 的距离,时间复杂度为$O(DNkk^{'} )$,本文例子为$10^{13}$ ,空间复杂度为$O(Nkk^{'})$ ,本文的例子为20G*4B=80GB。

索引结构

通过倒排索引,能大幅提高搜索效率

论文中提出使用 $kk^{'}$ 个倒排索引存储粗糙中心点 $c_i$对应的向量列表 $\mathcal{L}_i$ 。每个向量通过如下的格式表示,其中id是向量的索引id,code是对应的PQ中心点索引列表。PQ中每个组的中心点数量为$k^*$,则需要 $\lceil log_2{k^*} \rceil$ 个bit来表示哪一个中心点,共 $m*\lceil log_2{k^*} \rceil$个bit

当进行搜索时,可以通过 $q_c$ 函数获取对应簇下所有的向量

索引过程

  1. 通过量化器 $q_c$ 将向量 $y$ 映射到 $q_c(y)$
  2. 计算残差 $r(y) = y-q_c(y)$
  3. 将残差 $r(y)$ 量化到 $q_p(r(y))$ ,其中包含了m个分组
  4. 构造一个 $id|code$ 的 entry 并加入到 $q_c(y)$ 对应的倒排列表 $\mathcal{L}$
搜索过程

由于很多情况下,最近邻不一定当前的簇里,所以不仅要查找当前簇,还要查找邻近的簇。

  1. 计算 $\mathcal{C_c}$ 中与入参x最近的 $w$ 个中心点, $O((kk-w)*log{w})$
  2. 如果还有中心点没处理,取出一个中心点 $c_i$ ,并计算对应的 $r(x) = q_p(x-c_i)$。否则跳到步骤6
  3. 计算 $r(x)$ 与各个分组内的中心点距离,$O(m*\frac{D}{m}*k^{*}) = O(Dk^*)$
  4. 由于同一个倒排索引中对应的 $q_c(x)$ 与 $q_c(y)$ 是相同的,所以 x 与 y 的距离只要看残差距离 $d(r(x),r(y))$ 。由于 $r(x)$ 与各个中心点的距离都已经计算好,所以每个向量只需要查表m次即可。O(m)
  5. 返回步骤2
  6. 使用最小堆得到K个距离最小的向量,由于每个倒排索引预期元素数量为$\frac{N}{kk^{'}}$,所以耗时$O((\frac{N}{kk^{'}}-K) * logK)$

整个搜索过程的耗时为: $O((kk^{'}-w) * log{w}) + w * (O(Dk^{*})+O(m))+O((\frac{N}{kk^{'}}-K)*logK)$

实验效果

优化

  1. 残差量化中,不同的向量到各自中心的残差都放在一起进行量化了,其实隐含了不同聚类中的分布相同的假设。这个假设带来了一定误差,但不这么做的话内存占用就要扩大 $kk^{'}$ 倍了
  2. 对向量的不同分组方式会导致表现有很大的差异。论文中的实验显示,比起随机分组,相关的字段应该放在同一个分组在某些场景下可以使正确率提高2-3倍。在索引之前,可以通过一些相似度分析的方法将向量通过合适的顺序进行组织。
  3. $w$ 如果选取为1,会导致只查找当前簇内的向量,带来的结果可能比SDC还要差很多。作者在论文中的建议是取 $w=8$,但不同的场景下还是应该先进行测试以取得内存空间与耗时的平衡。

总结

Faiss 本质上是将向量编码为有限个向量的组合,将向量之间的距离计算转换为可提前计算的有限个向量之间的距离。

简化计算的关键:

  1. 候选向量缩小到邻域向量
  2. 乘积量化带来的表达能力跃升: $ka \ll (a)^k$
  3. 预先计算乘积量化结果的距离矩阵,使距离计算变为查表操作。(乘积量化使距离矩阵的空间需求在可忍受的范围内)

举个例子: 如下向量12、13

首先对他们进行粗糙聚类计算

发现他们的聚类id都是1

接着分成三组计算残差

对残差进行乘积量化

那么向量12、13的距离就可以直接通过累加三组乘积量化vector距离得到,其中vector距离都是提前计算好的

$d(12,13)=\sqrt{(a_{12})^2+(b_{23})^2+(c_{34})^2}$

参考文献: 《Product Quantization for Nearest Neighbor Search》: https://lear.inrialpes.fr/pubs/2011/JDS11/jegou_searching_with_quantization.pdf

About Me

2018.02至今 杭州嘉云数据 算法引擎

2017.6-2017.12 菜⻦网络-⼈工智能部-算法引擎

2016.09-2018.06 南京大学研究生

2015.07-2015.09 阿里巴巴-ICBU-实习

2012.09-2016.06 南京大学本科