Faiss 底层原理-来也科技

Faiss，Facebook的AI相似度搜索，是当下比较流行的高效相似度搜索框架，下面我们讲讲为什么是它最流行，以及在项目中怎么使用它：

Faiss是什么

在我们开始用faiss写代码之前，我们先来看看Faiss是什么：

首先Faiss是一个Facebook AI库，它可以使相似度的搜索变得容易；比如，我们有一个向量的集合，我们可以通过Faiss给他们建索引，然后用一个别的向量，我们可以用这个索引在这个集合中找到跟他相似的向量；
而且，他不只能够让我们构建向量索引、搜索向量，他可以通过参数的调节，使搜索变快，但是可能会导致召回或准确度下降，这也是这篇文章一个目的：让大家理解Faiss的底层的原理，更好地根据业务需求去调整参数

构造一些向量

第一件事情是我们需要一些数据，我们将从语义相似度测试hub仓库下载的多个的数据集，并进行解压、按字段进行分割转化为一个单一的列表：

import requestsfrom io import StringIOimport pandas as pdfrom sentence_transformers import SentenceTransformer

res = requests.get('https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/sick2014/SICK_train.txt')# create dataframedata = pd.read_csv(StringIO(res.text), sep='\t')data.head()

sentences = data['sentence_A'].tolist()sentences[:5]sentences = data['sentence_A'].tolist()sentence_b = data['sentence_B'].tolist()sentences.extend(sentence_b)  # merge themurls = [    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2012/MSRpar.train.tsv',    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2012/MSRpar.test.tsv',    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2012/OnWN.test.tsv',    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2013/OnWN.test.tsv',    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2014/OnWN.test.tsv',    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2014/images.test.tsv',    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2015/images.test.tsv']for url in urls:    res = requests.get(url)    # extract to dataframe    data = pd.read_csv(StringIO(res.text), sep='\t', header=None, error_bad_lines=False)    # add to columns 1 and 2 to sentences list    sentences.extend(data[1].tolist())    sentences.extend(data[2].tolist())

得到这个集合后，我们去重后可以得到1.45万个不同的句子。最后，我们通过sentence-BERT库对每个句子构建一个稀疏矩阵

sentences = [word for word in list(set(sentences)) if type(word) is str]model = SentenceTransformer('bert-base-nli-mean-tokens')sentence_embeddings = model.encode(sentences)

安装Faiss

如果你的操作系统是Linux，你非常幸运，因为Faiss在安装CUDA的Linux机器上默认支持内置的GPU优化；所以：

安装CUDA的Linux机器请用conda install -c pytorch faiss-gpu安装
其他的操作系统用conda install -c pytorch faiss-cpu；

当Faiss安装好之后，我们就可以用Python构建我们用简单易理解的IndexFlatL2来构建第一个向量索引。

IndexFlatL2

IndexFlatL2测量搜索向量与索引中所有向量的欧式距离，因此它非常简单，准确率也非常高，但是速度并不快。在Python中，我们可以初始化一个拥有跟我们句子向量一样多维度（768维）的IndexFlatL2索引，IndexFlatL2是不需要预训练的索引，我们可以直接将我们的向量加载到索引中。

import faissd = sentence_embeddings.shape[1]index = faiss.IndexFlatL2(d)index.add(sentence_embeddings)

然后我们用一个查询语句构造一个向量，用这个向量去查询跟他相似的k个邻居向量，他将返回对应向量在集合中的位置：

%%timek = 4xq = model.encode(["Someone sprints with a football"])D, I = index.search(xq, k)  # search#CPU times: user 27.9 ms, sys: 29.5 ms, total: 57.4 ms#Wall time: 28.9 ms[f'{i}: {sentences[i]}' for i in enumerate(I[0].tolist())]

4586 A group of football players is running in the field

10252 A group of people playing football is running past the person

12465 Two groups of people are playing football

190 A football player is running past an official

可以看出来，索引返回了跟我们查询向量xq最接近的前k个向量，而且返回的句子跟我们搜索的句子相似度非常匹配，都含有people running with a football或者在足球的上下文语义中。

如果我们需要进一步查看这些向量对应的值，我们可以通过下面的方式拿到：

vecs = np.zeros((k, d))# then iterate through each ID from I and add the reconstructed vector to our zero-arrayfor i, val in enumerate(I[0].tolist()):    vecs[i, :] = index.reconstruct(val)

速度

只是 IndexFlatL2 索引需要耗费大量的计算资源，而且它的伸缩性不是很好，当使用这种索引的时候，其实是遍历搜索，意味着：所有索引中的向量都将跟搜索向量xq进行比较，在我们的例子中，在每次搜索中将需要进行1.45万次欧式距离的计算。如果需要搜索的数据集到达100万，10亿，甚至更多，我们的索引将非常快地变得非常慢，甚至无法使用。所以我们需要其他类型的索引。

分区索引IndexIVFFlat

Faiss容许我们使用不同的方式去增加一些构建索引的步骤，从而优化我们的查询效率。一个非常流行的方式是将索引分割为不规则的泰森多边形，使用这种方式，我们首先确认查询向量xq所在的单元，再用上文讲到的IndexFlatL2在该单元中的向量集合中进行搜索，因此，我们可以减少搜索的范围，产出跟完全一致的结果比较接近的结果。

在代码实现上，我们首先用IndexFlatL2初始化一个索引，但这次我们只把这个L2索引当作量化器，然后我们把量化器传入分区索引IndexIVFFlat。在这里，我们增加了一个新的参数nlist，它用来指定有多少个分区；跟之前IndexFlatL2直接就可以得到索引，不需要分组转换所以不需要训练不一样的是：IndexIVFFlat使用了聚类，所以他需要提前训练；在训练后就可以跟之前一样，添加需要索引的数据。

nlist = 50  # how many cellsquantizer = faiss.IndexFlatL2(d)index = faiss.IndexIVFFlat(quantizer, d, nlist)index.train(sentence_embeddings)index.add(sentence_embeddings)

然后我们再用同样的xq来搜索：

%%timeD, I = index.search(xq, k)  # search#CPU times: user 3.83 ms, sys: 3.25 ms, total: 7.08 ms#Wall time: 2.15 ms

可以发现搜索时间明显下降，并且在这个例子里，这个近似搜索查询结果跟全量搜索没有不同；但经常会出现不一致的情况，这时如果IndexIVFFlat的近似搜索会返回次优解，我们可以通过增加查询范围来提高返回的准确率。在代码上，可以增加nprobe的值来实现。nprobe指定了搜索时需要搜索多少单元。

index.nprobe = 10%%timeD, I = index.search(xq, k)  # searchprint(I)#CPU times: user 5.29 ms, sys: 2.7 ms, total: 7.99 ms#Wall time: 1.54 ms

因为搜索的范围扩大了相应的搜索的相应时间也会增加。之间的向量的数量级及与不同nprobe响应关系图如下：

因

因为搜索的范围扩大了相应的搜索的相应时间也会增加。之间的向量的数量级及与不同nprobe响应关系图如下：

量化IndexIVFPQ

当向量的数据量持续增长，怎么对海量的数据构建向量索引，这时有另外一个关键的优化点：怎么从扁平的全索引，变成有压缩能力的产品量化PQ索引；PQ为什么会有压缩能力，是因为他跟IVF一样有个预处理，IVF的预处理是减少搜索距离，PQ的预处理是预处理两个向量的距离与相似度计算；PQ为了达到对相似度计算的预处理，它首先就会对向量进行压缩；整个过程分为3个步骤：

第一步：将原有向量分割成若干个子向量；
第二步：对每一组子向量进行聚类操作，找到每组子向量的多个中心点；
第三步：将这些子向量用离他最近的中心向量的ID代替，得到一个ID的向量；

在Faiss中，为实现这种索引，我们需要使用IndexIVFPQ索引，在加入向量数据之前，也需要对索引进行预训练；

m = 8  # number of centroid IDs in final compressed vectorsbits = 8 # number of bits in each centroidquantizer = faiss.IndexFlatL2(d)  # we keep the same L2 distance flat indexindex = faiss.IndexIVFPQ(quantizer, d, nlist, m, bits) index.train(sentence_embeddings)index.add(sentence_embeddings)

这时再来搜索：

index.nprobe = 10  # align to previous IndexIVFFlat nprobe value%%timeD, I = index.search(xq, k)#CPU times: user 3.04 ms, sys: 2.18 ms, total: 5.22 ms#Wall time: 1.33 ms

通过添加PQ，查询时间从之前的IVF ~7.5ms下降到~5ms，在这种小数据集上差别不大，但是随着向量数据量级的增加，响应时间将很快有明显的差别；尽管如此，但是我们在返回结果上也会有小小的差异，因为IVF和PQ对准确率都会有所损失，但是不是即使不是跟全量检索一样完美的结果，也能保证跟完美结果相差不远，而且因为预处理过，响应时间有数量级的提升：