海量文本求topk相似：faiss库初探

近似最近邻搜索的思想是：在海量数据的情况下，不必强求找到最相似的topk个数据，而只要搜索到相似度的确非常高的topk个数据即可，在可接受的范围内，牺牲精度以提高搜索的效率。

近似最近邻搜索一般可以实现两个目标：一是加快查询速度，二是降低内存压力。

那它是怎么实现的呢？

近似最近邻搜索中蕴含了降维和聚类的思想，以下是两种常用的实现方法：

一是基于哈希函数的局部敏感哈希，二是基于向量量化的乘积量化。

基于哈希的方法，是通过哈希函数把向量变换成二值码(如0101110101)，然后将向量距离近似成二值码距离（如海明距离），从而减少计算距离的时间。

局部敏感哈希就是一种基于哈希的方法，通过选择合适的哈希函数，将高维空间的数据映射到低维空间。

这起到降维的作用。

同时以非常大的概率，将在高维空间中相邻的点，映射到同一个桶（Bucket）中，而不相邻的点，映射到同一个桶的概率很小。

这起到聚类的作用。

在高维空间中相邻的数据，映射到低维空间后，依然保持相邻关系，这与传统的HashTable不同。

对数据库中的所有数据都进行哈希映射，把所有数据分到不同的桶里面。

过来一条查询数据，将其映射到某个桶内，然后和这个桶内的所有数据计算距离，从而得出topk相似的数据。

这就将在一个超大集合内查找相邻元素的问题，转化为了在一个很小的子集合内查找相邻元素的问题，计算量自然减少了。

这是什么意思？

向量用对应的类中心来近似，那这个类里面的所有向量不就是一样的吗？

错！

乘积量化（PQ，Product Quantization）是向量量化的代表方法，我们来看乘积量化是怎么做的，就能理解上面的话了。

上面说的聚类，可不是把所有的向量进行聚类，而是先把向量切分成若干个区域（类似BERT的Multi-Head的切分），再对每个区域的所有数据进行聚类。

比如数据有1万条，每条数据表示为100维的向量。

把100维向量切分为4个25维的子向量，第1个子向量集合仍然有1万个数据，第2、3、4个子向量集合也有1万个数据。

对于每个子向量集合（1万个数据）聚成K类，得到K个类中心。

每个子向量集合作为一个量化器，可以得到4个量化器（Quantizer）。

那么聚类中心一共是4×K个吗？

错！

聚类中心一共是 K^4 个。

因为这是乘积量化，而乘积指的是笛卡尔积。

假设集合A={a,b}，集合B={1,2,3,4}，也就是K=2，则两个集合的笛卡尔积为：

8个聚类中心，正是K的4次方。

那么向量用对应的聚类中心来近似，又是啥子？

比如一个向量被切分了4个子向量，4个子向量分别属于 a、b、a、b类。

那么该向量就表示为（a_1, b_2, a_3, b_4），也就是用聚类中心的id来表示。

从100维变成了4维，维度降低25倍。

如果这100维数据是float32类型的（32bits），压缩成8bits，那么压缩比为：

(100*32/4)/8=100

为啥你解释得这么细？

因为faiss这个库的乘积量化就是这么干的啊！

以上是前期的准备工作，可以看做是在训练一个模型，构建所有聚类中心的索引（Index）。

模型训练好了，来了一个查询向量，怎么搜索topk相似呢？

这就需要计算查询向量和数据库中某个向量的距离。

论文中有两种计算距离的近似方法：

对称方法（Symmetric distance computation ，SDC）。

非对称方法（Asymmetric distance computation ，ADC）。

对称方法就是指，将查询向量和被查询向量都用聚类中心的id来编码。

查询向量：（a_1, b_2, a_3, b_4）

被查向量：（a_1, a_2, b_3, a_4）

聚类中心两两之间的距离提前计算好了，假如为：

查询向量：（a_1, b_2, a_3, b_4）

被查向量：（a_1, a_2, b_3, b_4）

中心距离：（0,     1,     3,     0）

那么 sum(0,1,3,0)=4，就是查询向量和被查向量的距离。

对应的公式为：

非对称方法就是指，不需要将查询向量用聚类中心的id来编码，只对被查询向量编码。

将查询向量分为4个子向量后，直接求子向量和被查向量的聚类中心的距离，再加总求和。

查询向量：（1,     2,     3,     4）

被查向量：（a_1, b_2, a_3, b_4）

中心距离：（0,     1,     3,     0）

对应的公式为：

论文中选择的方法是非对称的方法。

用一句话来总结乘积量化：

乘积量化的本质是将高维向量空间分解为子空间的笛卡尔积，然后分别量化这些子空间。

四：faiss库的使用

说了一大堆，终于可以使用faiss了。

我们只安装CPU版就行，先安装好anaconda，再执行以下命令：

conda install faiss-cpu -c pytorch

官网说，安装CPU版需要一个库：BLAS。

我直接执行上面的命令也没报错，可能这个 BLAS 已经装好了。

04

faiss的乘积量化

亲测faiss中的乘积量化效果贼好，速度快而且省内存，所以我们直接先来看乘积量化的使用。

首先我们构造备查的向量数据和查询向量。

构造了200万个128维的备查向量，服从正态分布。

向量为float32类型，对应32bits（位），4bytes（字节）。

#coding:utf-8
import time
import numpy as np 

""" 
1: 向量128维，构造200万个向量，服从正态分布
"""
d = 128         
n_data = 2000000  
np.random.seed(0) 


data = []
mu = 3
sigma = 0.1

for i in range(n_data):
    data.append(np.random.normal(mu, sigma, d))

""" 向量为float32类型，即32bits（位数），4bytes（字节）"""
data = np.array(data).astype('float32')

再构造10个查询向量，同样是128维。

"""
2: 构造10个查询向量
"""

query = []
n_query = 10

for i in range(n_query):
    query.append(np.random.normal(mu, sigma, d))

query = np.array(query).astype('float32')

再接着构建乘积量化索引

"""
3: 构建索引
"""
import faiss

""" 将向量切分为8个子向量，每个向量16维，
所以m必须能被d整除"""
m = 8

""" 每个子向量集合都聚成100类 """
nlist = 100          

""" 取top10相似向量 """
k = 10

""" 构建量化器 """
start = time.time()
quantizer = faiss.IndexFlatL2(d) 

""" 每个子向量从32bits编码为8bits，
如果希望加大压缩比率，也可以编码为4bits """
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8)

index.train(data)
index.add(data)

""" 个人理解是计算查询向量和100个聚类中心的距离，然后取距离最小的10个类别，
然后再去这10个类别中找top10相似，类似于beam search，
因为nprobe=100时，时间消耗等同于没有聚类。
"""
index.nprobe = 10
end = time.time()
print("Time usage for train: {} seconds.".format(round((end - start),4)))

参数和方法介绍如下：

m ：按照上面乘积量化的介绍，需要将向量切分为多个子向量，我们选择切分为8个子向量。

显然子向量的个数m必须能被向量维度d整除。

nlist ：然后把每个子向量集合都聚成100类，也就是文档里面说的划分为100个维诺空间。

乘积量化索引应该用的是Kmeans来聚类。

k ：取topk相似的向量。

faiss.IndexIVFPQ ：倒排索引（IVF）+PQ编码。

倒排索引把向量用多个聚类中心的id来共同表示，以提高查询效率。

key为聚类中心的id，value为向量集合，如下：

但是索引中保存了原始向量（128维，32bits），在数据量大时会占用大量内存，甚至造成内存溢出。

于是使用PQ编码对原始向量进行有损压缩，从32bits压缩为8bits或4bits。

测试发现不使用乘积量化，构建好的索引占用了近1G内存，而使用乘积量化，构建好的索引只占用了32M内存。

index.nprobe ：在多少个维诺空间中查询。

个人理解是计算查询向量和100个聚类中心的距离，然后取距离最小的10个类别，然后再去这10个类别中找top10相似，类似于beam search。

nprobe越大，查询的精度越高，时间消耗越多，当nprobe=nlist时，时间消耗等同于没有做聚类的。

构建只花了14秒。

Time usage for train: 13.6571 seconds.

索引构建好之后，我们就可以查询topk相似的向量了。

首先我们通过查询在索引中已有的向量，来进行验证：

"""
4: 近似最近邻搜索
"""

""" 查询验证 """
start = time.time()
dis, ind = index.search(data[:10], k) 
end = time.time()

print(ind)

print("\nTime usage for search: {} seconds.\n".format(round((end - start),4)))

faiss支持并行查询，可以看到top1相似的向量就是向量本身，10条查询向量同时查询，只花了8ms。

然后进行真实查询，耗时也是8ms左右。

""" 真实查找 """

start = time.time()
dis, ind = index.search(query, k) 
end = time.time()

print(ind)

print("\nTime usage for search: {} seconds.\n".format(round((end - start),4)))

最后保存索引文件，文件大小只有32m。

"""
5: 保存索引和加载索引
"""
faiss.write_index(index, "index_IVFPQ.index")
index = faiss.read_index("index_IVFPQ.index")

05

faiss自定义id

"""
6: 自定义向量id
"""

""" id必须为int类型 """
ids = np.arange(2000000,4000000)

quantizer = faiss.IndexFlatL2(d) 
index_ = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8)

index_.train(data)

""" 自定义id """
index_.add_with_ids(data,ids)

dis, ind = index_.search(query, k) 

print(ind)

"""
7：索引工厂
"""
index_fac = faiss.index_factory(d,"OPQ16_64,IMI2x8,PQ16") 

start = time.time()
index_fac.train(data)
index_fac.add(data)
end = time.time()
print("\nTime usage for train: {} seconds.\n".format(round((end - start),4)))

start = time.time()
dis, ind = index_fac.search(data[:10], 10)  # 真实查询
end = time.time()

print(ind)
print("\nTime usage for search: {} seconds.\n".format(round((end - start),4)))

# indexes that are useful when there is no limitation on memory usage
unlimited_mem_keys = [
"IMI2x10,Flat", "IMI2x11,Flat",
"IVF4096,Flat", "IVF16384,Flat",
"PCA64,IMI2x10,Flat"]

# memory limited to 16 bytes / vector
keys_mem_16 = [
'IMI2x10,PQ16', 'IVF4096,PQ16',
'IMI2x10,PQ8+8', 'OPQ16_64,IMI2x10,PQ16'
    ]

# limited to 32 bytes / vector
keys_mem_32 = [
'IMI2x10,PQ32', 'IVF4096,PQ32', 'IVF16384,PQ32',
'IMI2x10,PQ16+16',
'OPQ32,IVF4096,PQ32', 'IVF4096,PQ16+16', 'OPQ16,IMI2x10,PQ16+16'
    ]

"""
8：精确查找作为baseline，进行对比
"""

index_exact = faiss.IndexFlatL2(d)  

""" 不需要训练 """
index_exact.add(data)

start = time.time()
dis, ind = index_exact.search(data[:10], 10)
end = time.time()

print(ind)
print("\nTime usage for search: {} seconds.\n".format(round((end - start),4)))

参考资料：

1：《一文尽览近似最近邻搜索中的哈希与量化方法》