Faiss原理和使用

参考自https://github.com/facebookresearch/faiss/wiki，https://blog.csdn.net/Kaiyuan_sjtu/article/details/121551473

Faiss

Faiss是一个用于高效相似性搜索和密集向量聚类的库。它包含搜索任意大小的向量集（大小由RAM决定）的算法。它还包含用于评估和参数调整的支持代码。Faiss用C++编写，带有完整的Python包装器。一些最有用的算法是在GPU上实现的。

什么是相似性搜索？

给定一组维度为$d$的向量 { x 1 , . . . , x n } \{x_1,...,x_n\} {x1​,...,xn​}，Faiss会在RAM中构建一个数据结构。构建结构之后，当给定一个维度为$d$的新向量$x$时，它会有效地执行以下操作：
$$i=argmi{n}_i||x-x_i||$$
其中，$||\cdot ||$是欧几里得距离(L2)。

在Faiss术语中，数据结构是一个索引，一个具有add方法以添加x_i向量的对象。注意，x_i的维度被假定为固定的。

这就是Faiss的全部内容。它还可以：

• 不仅返回最近邻居，还返回第二近邻、第三近邻、…、第k近邻
• 一次搜索多个向量而不是一个（批处理）。对于许多索引类型，这比逐个搜索向量更快
• 以精度换取速度，即使用速度快10倍或使用内存少10倍的方法，代价是10%的时间会给出错误结果
• 执行最大内积搜索$argma{x}_i(x,x_i)$而不是最小欧几里得搜索。对其他距离（$L1$， $Linf$等）也有有限的支持
• 返回查询点给定半径内的所有元素（范围搜索）
• 将索引存储在磁盘上而不是RAM中
• 索引二进制向量而不是浮点向量
• 根据向量ID上的谓词忽略索引向量的子集

Faiss的原理

对于m个向量和n个向量做相似度比较，用暴力搜索的时间复杂度就是O(mn)，在向量个数非常多的情况下，这种资源消耗十分巨大的。Faiss使用索引，并且以一些精度换取速度，能够大大加快搜索速度。

Faiss的核心原理有两个部分：

Product Quantizer（PQ——乘积量化）

矢量量化方法，即vector quantization，其具体定义是将向量中的点用一个有限子集来进行编码的过程。常见的聚类算法都是一种矢量量化方法。而在ANN（Approximate Nearest Neighbor，近似最近邻）搜索问题中，向量量化方法又以乘积量化（PQ，Product Quantization）最为典型。

PQ的Pretrain

PQ的Pre-train过程分为两步操作，第一步Cluster，第二步Assign，这两步合起来就是Faiss数据流的Train阶段，以一个128维的向量库为例：
向量库向量个数为N，则共有$N\ast 128$维，那么将其进行切分，比如将一个向量切成4个子段，那么就是$N\ast 4\ast 32$维。这个切分的参数为$M$，$M=4$代表每一个向量被切成了4段。

在切分后，按所有向量的每一段进行Clustering，比如所有向量的第一段取出来做Clustering得到256个簇心（经验值），这样总共得到$256\ast M$个簇心。

做完Cluser后，对所有向量进行Assign操作，Assign就是把原来的N维的向量映射到M维，比如上面的$128，4$，对于每一段向量，都可以找到对应的最近的簇心ID，这样一个128维的向量可以由4个簇心ID表示，这样就完成了向量的压缩，把$N\ast 128$压缩成了$N\ast 4$。

PQ的查询

对于查询向量，以相同的方式把其分成$M$段，计算每一段向量与之前预训练好的簇心的距离，相当于本来是128维的查询向量和$N\ast 128$维的向量库进行查询，现在变成了4*32的查询向量和$N\ast 4$维的向量库进行查询，当N变大时，速度差距越来越大。
注意：相当于将向量的原始表示转化成向量的簇心ID表示，即一个向量[x1,x2,x3…]转换成了[y1,y2,y3,y4]，其中y1,y2,y3,y4对应的是簇心ID。

Inverted File System（IVFPQ——倒排乘积量化）

PQ优化了向量距离计算的过程，但是对于每个库本身，仍然需要遍历整个库，因为对于查询向量，还是需要和库的全部簇心进行计算并相加，效率依旧不够高，所以就有了Faiss用到的另外一个关键技术——Inverted File System。

如果能够通过某种手段快速将全局遍历锁定为感兴趣区域，则可以舍去不必要的全局计算以及排序。倒排PQ乘积量化的“倒排”，就是这种思想，通过聚类的方式实现感兴趣区域的快速定位。

要想减少需要计算的目标向量的个数，做法就是直接对库里的所有向量做KMeans CLustering，假设簇心个数为1024，那么每来一个query向量，首先计算其与1024个粗聚类簇心的距离，然后选择距离最近的top N个簇，只计算查询向量与这几个簇底下的向量的距离，计算距离的方式就是前面的PQ。

Faiss具体的实现有个小细节，在计算查询向量和簇底下的向量的距离的时候，所有向量都会被转化成与簇心的残差，这类似于归一化操作，使得后面用PQ计算距离更准确一点。使用IVF后，需要计算的向量个数就少了几个数量级，提高向量检索的速度。

Faiss的使用

• Flat

import numpy as np

# 定义数据维度
d = 64  
# 数据库中数据点的数量
nb = 100000  
# 查询的数量
nq = 10000  
# 设置随机种子以保证结果可复现
np.random.seed(1234)  
# 生成数据库中的随机数据点
xb = np.random.random((nb, d)).astype('float32')
# 为每个数据点的第一维添加一个唯一的小数值
xb[:, 0] += np.arange(nb) / 1000.
# 生成查询点的随机数据
xq = np.random.random((nq, d)).astype('float32')
# 为每个查询点的第一维添加一个唯一的小数值
xq[:, 0] += np.arange(nq) / 1000.

# 导入 Faiss 库
import faiss  
# 创建一个 L2 距离的 Flat 索引，维度为 d
index = faiss.IndexFlatL2(d)  
# 打印索引是否已经训练（Flat 索引不需要训练）
print(index.is_trained)  
# 将数据库中的向量添加到索引中
index.add(xb)  
# 打印索引中的向量总数
print(index.ntotal)

# 设置要检索的最近邻个数
k = 4  
# 对数据库中的前 5 个点进行搜索，检索每个点的 4 个最近邻
D, I = index.search(xb[:5], k)  
# 打印最近邻的索引
print(I)  
# 打印最近邻的距离
print(D)  
# 对所有查询点进行搜索，检索每个查询的 4 个最近邻
D, I = index.search(xq, k)  
# 打印前 5 个查询点的最近邻索引
print(I[:5])  
# 打印最后 5 个查询点的最近邻索引
print(I[-5:])  

这段代码演示了如何使用 Faiss 库创建一个基于 L2 距离（欧几里得距离）的最近邻搜索索引，向索引中添加数据点，然后使用这些数据点进行查询以找到最接近的 k 个邻居。代码中的 search 函数返回两个数组：D 包含与最近邻点的距离，而 I 包含最近邻点在数据库中的索引。

• IVFFlat

import numpy as np

# 定义数据的维度
d = 64  
# 数据库中数据点的数量
nb = 100000  
# 查询的数量
nq = 10000  
# 设置随机种子以保证结果可复现
np.random.seed(1234)  
# 生成数据库中的随机数据点
xb = np.random.random((nb, d)).astype('float32')
# 为每个数据点的第一维添加一个唯一的小数值
xb[:, 0] += np.arange(nb) / 1000.
# 生成查询点的随机数据
xq = np.random.random((nq, d)).astype('float32')
# 为每个查询点的第一维添加一个唯一的小数值
xq[:, 0] += np.arange(nq) / 1000.

import faiss

# 定义量化列表的大小
nlist = 100  
# 定义要检索的最近邻个数
k = 4  
# 创建量化器，使用 L2 距离的 Flat 索引
quantizer = faiss.IndexFlatL2(d)  
# 创建倒排文件索引（IVF）结合量化器
index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
# 指定使用 L2 距离度量，如果不指定，默认执行内积搜索

# 确保索引尚未训练
assert not index.is_trained  
# 训练索引，使用整个数据库
index.train(xb)  
# 确保索引已经训练
assert index.is_trained  

# 将数据库向量添加到索引中
index.add(xb)                  
# 执行实际的搜索，检索每个查询的 4 个最近邻
D, I = index.search(xq, k)     
# 打印最后 5 个查询点的最近邻索引
print(I[-5:])                  

# 设置 nprobe 参数，nprobe 默认是 1，尝试设置为更大的值
index.nprobe = 10              
# 使用更新的 nprobe 参数重新执行搜索
D, I = index.search(xq, k)     
# 再次打印最后 5 个查询点的最近邻索引
print(I[-5:])                  

这段代码演示了如何使用 Faiss 库创建一个倒排文件索引（Inverted File Index，IVF）结合量化的索引结构，用于高效的相似性搜索。IndexIVFFlat 索引首先使用量化器将数据库划分为多个列表（由 nlist 参数指定），然后对每个列表内的向量进行最近邻搜索。nprobe 参数控制搜索时考虑的列表数量，增加 nprobe 的值可以提高搜索的精度，但可能会降低搜索速度。

• IVFPQ

import numpy as np

# 定义数据的维度
d = 64                           # dimension
# 数据库中数据点的数量
nb = 100000                      # database size
# 查询的数量
nq = 10000                       # nb of queries
# 设置随机种子以保证结果可复现
np.random.seed(1234)             
# 生成数据库中的随机数据点
xb = np.random.random((nb, d)).astype('float32')
# 为每个数据点的第一维添加一个唯一的小数值
xb[:, 0] += np.arange(nb) / 1000.
# 生成查询点的随机数据
xq = np.random.random((nq, d)).astype('float32')
# 为每个查询点的第一维添加一个唯一的小数值
xq[:, 0] += np.arange(nq) / 1000.

import faiss

# 定义每个列表中的聚类中心数
nlist = 100
# 定义每个聚类中的码本大小
m = 8
# 定义要检索的最近邻个数
k = 4  
# 创建量化器，使用 L2 距离的 Flat 索引
quantizer = faiss.IndexFlatL2(d)  
# 创建一个使用乘积量化（PQ）的倒排文件索引
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8)
                                  # 8 指每个子向量被编码为 8 位
# 训练索引，使用整个数据库
index.train(xb)
# 将数据库向量添加到索引中
index.add(xb)
# 对数据库中的前 5 个点进行搜索，检索每个点的 4 个最近邻进行测试
D, I = index.search(xb[:5], k) 
# 打印最近邻的索引
print(I)
# 打印最近邻的距离
print(D)
# 设置 nprobe 参数，使结果与之前的实验具有可比性
index.nprobe = 10              
# 使用更新的 nprobe 参数执行搜索
D, I = index.search(xq, k)     
# 打印最后 5 个查询点的最近邻索引
print(I[-5:])

这段代码演示了如何使用 Faiss 库创建一个使用乘积量化（Product Quantization, PQ）的倒排文件索引（Inverted File Index, IVF）。乘积量化是一种压缩技术，可以减少存储需求并加速最近邻搜索，同时保持较高的精度。IndexIVFPQ 索引首先使用量化器将数据空间划分为多个聚类（由 nlist 参数指定），然后在每个聚类内部使用乘积量化技术对向量进行编码和搜索。m 参数定义了每个聚类中的码本大小，而最后一个参数 8 指定了每个子向量使用 8 位进行编码。nprobe 参数控制搜索时考虑的聚类数量，增加 nprobe 的值可以提高搜索的精度，但可能会降低搜索速度。

• GPU、Multipe-GPUs、PQFastScan、PQFastScanRefine、RefineComparison（https://github.com/facebookresearch/faiss/tree/main/tutorial/python）