Faiss为啥这么快?原来是量化器在做怪!1:https://developer.aliyun.com/article/1507793
2.乘积量化(Product Quantization)
PQ是一种用于高维向量压缩和相似性搜索的技术,特别适用于大规模向量数据集的近似最近邻搜索。PQ 技术将高维向量分解为多个子空间,并对每个子空间进行独立的向量量化,从而实现高效的向量压缩和搜索。下面将对量化和检索两个方面进行介绍:
PQ的论文地址:
https://arxiv.org/pdf/2401.08281.pdf
(1)PQ量化
量化过程有如下3个步骤:
1. 向量分解:将原始高维向量分解为多个子向量,每个子向量属于一个子空间。
2. 子空间量化:对每个子空间进行独立的向量量化,将连续的向量空间划分为离散的子空间。
3. 编码:将原始向量编码为子空间的离散码字,以表示原始向量在每个子空间中的位置。
下面来看一个例子:
数据集是一个1000x1024的矩阵,每个向量维度1024:
向量分解:先将每个1024维的向量平均分成8个128维的子向量,如下图所示:
子空间量化:然后对这8组子向量分别使用k-means聚成256类。下图中(1)竖着的一列为一组,每组进行聚类,聚类成256个簇,形成8x256的码表(2)。然后将每个128维的原向量转换成簇的中心点,如下图中的(3)
编码:接下来将每一个格中的128维向量根据(2)量化成一个簇ID(0-256 占8bit),这样原始1024维浮点数(32位)向量便压缩成8个8位整数,即下图(4)中粉色矩阵中的一行。
在上面的例子中,子向量数量和每个子向量簇的数量是两个超参数,经过实验子向量数量选择8和簇的数量选择256通常是最佳的。
(2)PQ检索
检索过程非常类似于找回+排序:
1. 量化查询向量:使用上面的压缩过程将查询向量转换成中心点ID矩阵。
2.召回:将查询的 PQ 编码与码本中的所有 PQ 编码进行比较,寻找与查询编码最接近的候选向量,比如计算海明距离。
3.排序:对于每个候选编码,找到对应的原始向量,然后精确计算相似度,进行排序。
这个流程大致如下:
上图中M=8,nlist=8
Faiss中PQ相关的有如下相关算法:
代码示例如下:
""" 乘积量化示例 """ import numpy as np import faiss import time def test_pq(): # 设定量化器参数 quantizer = faiss.ProductQuantizer(d, M, nlist) # 对示例数据进行训练 quantizer.train(data) # 进行量化 codes = quantizer.compute_codes(data) print("data0:", data[0]) print("codes0:", codes[0]) def test_index(): # 向索引中添加数据 t = time.time() index.train(data) index.add(data) cost1 = time.time() - t # 进行近似最近邻搜索 k = 50 # 返回最近邻的数量 t = time.time() D, I = index.search(query, k) cost2 = time.time() - t return D, I, cost1, cost2 if __name__ == '__main__': d = 1024 M = 8 nbits = 8 nlist = 500 nbits_per_idx = 4 n = 100000 np.random.seed(0) # 生成一些随机向量作为示例数据 data = np.random.rand(n, d).astype(np.float32) # 定义查询向量 query = np.random.rand(1, d).astype(np.float32) # 量化器的示例 # test_pq() # 暴力检索的示例 index = faiss.IndexFlat(d, faiss.METRIC_L2) D, I, cost1, cost2 = test_index() print("Flat索引 最近邻的距离:", D) print("Flat索引 最近邻的索引:", I) print("Flat索引 耗时:", cost1, cost2) # PQ索引的示例 index = faiss.IndexPQ(d, M, nbits, faiss.METRIC_L2) D, I, cost1, cost2 = test_index() print("PQ索引 最近邻的距离:", D) print("PQ索引 最近邻的索引:", I) print("PQ索引 耗时:", cost1, cost2) # IVF+PQ的示例 # quantizer = faiss.IndexPQ(d, M, nbits, faiss.METRIC_L2) quantizer = faiss.IndexFlat(d, faiss.METRIC_L2) index = faiss.IndexIVFPQ(quantizer, d, nlist, M, nbits_per_idx, faiss.METRIC_L2) D, I, cost1, cost2 = test_index() print("IVF+PQ索引 最近邻的距离:", D) print("IVF+PQ索引 最近邻的索引:", I) print("IVF+PQ索引 耗时:", cost1, cost2)
3.加法量化
加法量化(Additive Quantization)的基本思想是将原始向量表示为多个部分的和,每个部分都可以独立进行量化。在加法量化中,每个部分由一个码本表示,最终的量化结果是这些部分码本的加和。
Faiss中AQ相关的算法有ResidualQuantizer、LocalSearchQuantizer和AdditiveQuantizer,其中AdditiveQuantizer是ResidualQuantizer和LocalSearchQuantizer的父类。
(1)残差量化(Residual Quantizer)
残差量化,涉及对数据进行多次量化,每次量化(通常使用的是简单的标量量化)都使用相同的量化级别,但重建过程中会考虑前一次量化产生的误差。这种方法的关键在于,通过只处理量化过程产生的误差,可以更有效地压缩数据,同时尽量减少信息损失。
过程包括以下步骤:
a.初步量化:对数据进行一次简单的量化(K-means),将其划分到有限数量的级别中。
b.计算残差:计算初步量化后的数据与原始数据之间的差异,这些差异被称为“残差”。
c.再次量化:对计算出的残差进行再次量化,这次量化可以更加精确,因为残差通常比原始数据小很多。
d.迭代过程:b、c两个步骤迭代进行,即对残差的量化结果再次计算残差,并进行量化。
Faiss中残差量化相关的有如下相关算法:
代码示例:
""" 残差量化示例 """ import numpy as np import faiss import time def test_rq(): # 设定量化器参数 quantizer = faiss.ResidualQuantizer(d, M, nbits) # 对示例数据进行训练 quantizer.train(data) # 进行量化 codes = quantizer.compute_codes(data) print("data0:", data[0]) print('shape:', codes.shape) print("codes0:", codes[0]) def test_index(): # 向索引中添加数据 t = time.time() index.train(data) index.add(data) cost1 = time.time() - t # 进行近似最近邻搜索 k = 50 # 返回最近邻的数量 t = time.time() D, I = index.search(query, k) cost2 = time.time() - t return D, I, cost1, cost2 if __name__ == '__main__': d = 1024 M = 3 nbits = 8 nlist = 100 n = 100000 np.random.seed(0) # 生成一些随机向量作为示例数据 data = np.random.rand(n, d).astype(np.float32) # 定义查询向量 query = np.random.rand(1, d).astype(np.float32) # 量化器的示例 test_rq() # 暴力检索的示例 index = faiss.IndexFlat(d, faiss.METRIC_L2) D, I, cost1, cost2 = test_index() print("Flat索引 最近邻的距离:", D) print("Flat索引 最近邻的索引:", I) print("Flat索引 耗时:", cost1, cost2) # SQ索引的示例 index = faiss.IndexResidualQuantizer(d, M, nbits, faiss.METRIC_L2) D, I, cost1, cost2 = test_index() print("RQ索引 最近邻的距离:", D) print("RQ索引 最近邻的索引:", I) print("RQ索引 耗时:", cost1, cost2) # IVF+SQ的示例 # quantizer = faiss.IndexResidualQuantizer(d, M, nlist, faiss.METRIC_L2) quantizer = faiss.IndexFlat(d, faiss.METRIC_L2) index = faiss.IndexIVFResidualQuantizer(quantizer, d, nlist, M, nbits, faiss.METRIC_L2, faiss.ResidualQuantizer.ST_norm_qint8) D, I, cost1, cost2 = test_index() print("IVF+SQ索引 最近邻的距离:", D) print("IVF+SQ索引 最近邻的索引:", I) print("IVF+SQ索引 耗时:", cost1, cost2)
(2)局部检索量化(LocalSearchQuantizer)
局部检索量化,先使用k-means将原始高维向量空间划分为若干个簇,然后将原始的高维向量其映射到最近的簇心,最后再使用量化方法将高维向量转换成一个代表性的低维向量。检索的时候使用局部搜索的方法在这些低维向量上进行最近邻搜索。
LocalSearchQuantizer是下面两篇论文的实现:
Revisiting additive quantization Julieta Martinez, et al. ECCV 2016
LSQ++: Lower running time and higher recall in multi-codebook quantization Julieta Martinez, et al. ECCV 2018
示例代码如下:
""" 局部检索量化示例 """ import numpy as np import faiss import time def test_lsq(): # 设定量化器参数 quantizer = faiss.LocalSearchQuantizer(d, M, nbits) # 对示例数据进行训练 quantizer.train(data) # 进行量化 codes = quantizer.compute_codes(data) print("data0:", data[0]) print('shape:', codes.shape) print("codes0:", codes[0]) def test_index(): # 向索引中添加数据 t = time.time() index.train(data) index.add(data) cost1 = time.time() - t # 进行近似最近邻搜索 k = 50 # 返回最近邻的数量 t = time.time() D, I = index.search(query, k) cost2 = time.time() - t return D, I, cost1, cost2 if __name__ == '__main__': d = 1024 M = 3 nbits = 8 nlist = 100 n = 10000 np.random.seed(0) # 生成一些随机向量作为示例数据 data = np.random.rand(n, d).astype(np.float32) # 定义查询向量 query = np.random.rand(1, d).astype(np.float32) # 量化器的示例 test_lsq() # 暴力检索的示例 index = faiss.IndexFlat(d, faiss.METRIC_L2) D, I, cost1, cost2 = test_index() print("Flat索引 最近邻的距离:", D) print("Flat索引 最近邻的索引:", I) print("Flat索引 耗时:", cost1, cost2) # LSQ索引的示例 index = faiss.IndexLocalSearchQuantizer(d, M, nbits, faiss.METRIC_L2) D, I, cost1, cost2 = test_index() print("LSQ索引 最近邻的距离:", D) print("LSQ索引 最近邻的索引:", I) print("LSQ索引 耗时:", cost1, cost2) # IVF+LSQ的示例 quantizer = faiss.IndexFlat(d, faiss.METRIC_L2) index = faiss.IndexIVFLocalSearchQuantizer(quantizer, d, nlist, M, nbits, faiss.METRIC_L2, faiss.ResidualQuantizer.ST_norm_qint8) D, I, cost1, cost2 = test_index() print("IVF+LSQ索引 最近邻的距离:", D) print("IVF+LSQ索引 最近邻的索引:", I) print("IVF+LSQ索引 耗时:", cost1, cost2)
(3)加法量化(Additive Quantizer)
Faiss中的AdditiveQuantizer是将残差量化或者局部量化的结果加起来作为向量的量化结果,不过用的不是很多,就不介绍了。
Faiss中还有一些量化器,是上面这些量化器的排列组合,看名字就知道功能,使用频率也不是很高:
类名 | 描述 |
faiss.AdditiveQuantizer | 加法量化 |
faiss.AdditiveQuantizerFastScan | 使用了FastScan的加法量化 |
faiss.ProductLocalSearchQuantizer | 乘积量化+局部检索量化 |
faiss.ProductResidualQuantizer | 乘积量化+残差量化 |
faiss.ProductResidualQuantizerFastScan使用了FastScan的乘积量化+残差量化faiss.ProductLocalSearchQuantizerFastScan使用了FastScan的乘积量化+局部检索量化
最后一句话总结,如果向量维度不是很大(1024以内),数据量也不是很多(100w以内),内存允许的情况下可以使用Flat暴力搜索,效率其实可以接受;如果维度和数据量都很大,还是老老实实用量化吧,比如IndexIVFPQ还是很香的。
Faiss中的量化器就介绍到这里,欢迎点赞、收藏,关注不迷路(*^__^*)