Faiss,Facebook的AI相似度搜索,是当下比较流行的高效相似度搜索框架,下面我们讲讲为什么是它最流行,以及在项目中怎么使用它:
Faiss是什么
在我们开始用faiss写代码之前,我们先来看看Faiss是什么:
  1. 首先Faiss是一个Facebook AI库,它可以使相似度的搜索变得容易;比如,我们有一个向量的集合,我们可以通过Faiss给他们建索引,然后用一个别的向量,我们可以用这个索引在这个集合中找到跟他相似的向量;
  2. 而且,他不只能够让我们构建向量索引、搜索向量,他可以通过参数的调节,使搜索变快,但是可能会导致召回或准确度下降,这也是这篇文章一个目的:让大家理解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对准确率都会有所损失,但是不是即使不是跟全量检索一样完美的结果,也能保证跟完美结果相差不远,而且因为预处理过,响应时间有数量级的提升:



图片

总结:
通过这篇文章,我们对用Faiss构建高效查询索引的几种方式有了全面的介绍,也通过这些介绍,我们对每种索引怎么选择不同的参数去满足不同业务场景准确率/速度的要求有了更好的理解。总的来说这些让人印象深刻的能力都是基于Faiss是一个非常简单易用的高效向量搜索引擎,感谢Faiss!



原文链接:https://jamescalam.medium.com/free-course-on-vector-similarity-search-and-faiss-9b3e91a91384

翻译者:谭繁华