Faiss库系列2_原理分析

Faiss库系列2_原理分析

1 引入

  Faiss库为向量近邻搜索提供精确的暴力匹配方法IndexFlatL2和IndexFlatIP，前者使用L2距离衡量向量相似度，而后者采用点积衡量。暴力匹配存在的最大问题是耗时严重，针对这个问题，Faiss库利用泰森多边形（voronoi图）将搜索空间分为数个小单元（cell，划分的数量可有超参数nlist控制），并在搜索时控制搜索的单元数量（由nprobe控制）；经过这样的处理后，近邻搜索的速度得到了大幅提升，如IndexIVFFlat等索引方法。除了搜索耗时的问题，当搜索的向量达到10亿级别，内存的消耗将非常严重，于是Faiss库采用了乘积量化（Product Quantizer，简称PQ）和主成分分析（PCA）降维寻求搜索效率与内存之间的平衡，如IndexIVFPQ等方法。最后，Faiss库为了进一步加速近邻搜索效率，对多个核心算法进行GPU实现，支持单/多GPU以及CPU与GPU之间的无缝连接。
  由以上Faiss库提供的接口逻辑可以知道，Faiss库的三大核心算法是K-means聚类/K-selection、PCA以及PQ编码/解码。针对这三个核心算法，Faiss库提供了对应的接口，分别为：faiss.Kmeans()、faiss.PCAMatrix()以及faiss.ProductQuantizer()。接下来的章节将对Faiss库相关的内容进行简要介绍。

2 近邻搜索问题

  Faiss库关注向量集合的相似搜索问题。形式化定义为：给定查询向量$x\in {\mathbb{R}}^d$以及向量集合$[y_i{]}_{i=0:l}(y_i\in {\mathbb{R}}^d)$，按某种度量进行相似向量搜索，此处以L2距离为例：
$$\mathcal{L}=k-argmi{n}_{i=0:l}||x-y_i|{|}_2$$
  如果仅查询出前k个最相似的向量，被称为K-selection。形式化定义为：从向量集合$[y_i{]}_{i=0:l}(y_i\in {\mathbb{R}}^d)$查询出k个向量$[a_{s_i}{]}_{i=0:k}$，满足关系$a_{s_i}\le a_{s_{i+1}}$，其中$0\le [s_i{]}_{i=0:k}\le l$。

3 批量查询

  批量查询也是Faiss库所采用的查询模式。形式化定义为：给定$n_q$个查询向量$[x_j{]}_{j=0:n_q}(x_j\in {\mathbb{R}}^d)$，那么对于K-selection将返回$n_q\times k$的矩阵。在具有多CPU线程或GPU时，批量查询有利于并行化的查询处理。

4 暴力搜索（精确搜索）

  精确解需要计算查询向量到所有向量的距离，即需要计算距离矩阵$\mathcal{D}=[||x_j-y_i|{|}_2^2{]}_{j=0:n_1,i=0:l}\in {\mathbb{R}}^{n_q\times l}$。在工程实现中，需对距离度量进行如下分解：
$$||x_j-y_i|{|}_2^2=||x_j|{|}^2+||y_i|{|}^2-2⟨x_j,y_i⟩$$
  由以上分解可以看出，前两项可以由矩阵$X、Y$提前计算出来，而第三项是限制距离计算的瓶颈（对应于$X{Y}^T$）。

4 压缩域搜索（乘积量化）

  为了降低向量查询的搜索时间，压缩域搜索关注于近似近邻搜索。Faiss库采用论文[Product quantization for nearest neighbor search]中的IVFADC索引结构。现介绍如下：
  IVFADC索引依于两层量化子对向量集进行编码，对向量集中的向量y将表示为如下形式：
$$y\cong q(y)=q_1(y)+q_2(y-q_1(y))$$
  其中$q_1:{\mathbb{R}}^d\to {\mathcal{C}}_1\in {\mathbb{R}}^d$、$q_2:{\mathbb{R}}^d\to {\mathcal{C}}_2\in {\mathbb{R}}^d$均为量化子（例如输出为有限元素的函数）；由以上公式可知，$q(y)$可由$q_1(y)与q_2(y-q_1(y)$共同编码。其中，第一层编码$q_1(y)$为粗量化，第二层$q_2(y-q_1(y)$为对第一层量化残差的精细量化。
  定义非对称距离计算（ADC，Asymmetric Distance Computation）获取查询的近似结果，详细形式如下：
$$L_{ADC}=k-argmi{n}_{i=0:l}||x-q(y_i)|{|}_2$$
  由量化关系可知，IVFADC搜索为非精确查询。同时由量化的层次性，查询的预选择由第一层编码$q_1(y)$决定，可形式化表示为如下：
$$L_{IVF}=\tau -argmi{n}_{c\in {\mathcal{C}}_1}||x-c|{|}_2$$
  其中的参数$\tau$被称为multi-probe parameter。然后，在第一层编码$q_1(y)$预选择的向量集中进行第二层的向量搜索，具体形式如下：
$$L_{IVFIDC}=k-argmi{n}_{i=0:ls.t.q_1(y_i)\in L_{IVF}}||x-q(y_i)|{|}_2$$
  由此可见，IVFADC进行距离计算主要依赖与第二层$q_2(y-q_1(y)$的向量搜索，其需要进行距离计算的范围减小了，所以提升了查询的效率。