skyitachi's blogIVFPQ的原理 2024-08-04

主要内容

  • 介绍IVF和PQ的简单原理
  • 介绍IVFPQ的主要流程

Kmeans 聚类算法

在IVF和PQ中都涉及到聚类算法，简单介绍下kmeans聚类算法原理

1
2
3
4
1. 随机选择k个初始中心点
2. 将每个向量分配到最近的中心点所在的簇
3. 重新计算每个簇的中心点（找出质心）
4. 重复步骤2和3，直到簇中心不再变化或者达到迭代次数

IVF的与原理介绍

Inverted File (IVF) 是一种索引结构，通常用于信息检索系统，特别是在全文搜索引擎中

IVF的原理简单直接:

向量中的IVF索引就是将样本向量首先生成nlist个聚类，然后每次向量查询的时候只要比较查询向量和nlist个聚类中心的距离，可以找出nprobe个最近的聚类中心，然后在每个聚类里暴力搜索出距离最近的K个向量

代码示例

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import numpy as np
import faiss

# 生成一些示例数据
d = 64  # 向量维度
nb = 100000  # 数据库大小
nq = 10  # 查询数量
np.random.seed(1234)  # 确保可重复性
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')

# 构建 IVF 索引
nlist = 1024 # 聚类数量
quantizer = faiss.IndexFlatL2(d)  # 量化器
index = faiss.IndexIVFFlat(quantizer, d, nlist)

# 训练索引
index.train(xb)
index.add(xb)

index.nprobe = 256

# 查询
k = 4  # 返回的最近邻数量
D, I = index.search(xq, k)  # 搜索

# 打印结果
print("查询结果 (距离, 索引):")
for i in range(nq):
    print(f"查询 {i}

PQ(Product Quantization)的原理与介绍

PQ 主要用于向量的压缩, 一般而言向量都是高维数据, 假设一个1024维的float32类型的向量，单个向量的大小就是1024 * 4 = 4kb, 使用PQ压缩后往往能极大缩小单个向量的大小，当然不可避免的会丢失精度。

PQ向量压缩的主要步骤

1
2
1. 将d维向量分成M组，在每一组内找到ksub个聚类中心
2. 将每一组的向量用离它最近聚类中心的id来表示，那么每组内的向量的大小就是nbits（2^nbits=ksub）, 整体压缩过后的向量大小就是M * nbits, 一般而言M选择8，nbits一般是8或16

搜索压缩过后的样本向量

1
2
3
1. 将查询向量分为M组，每一组内和ksub个聚类中心计算距离, 并使用一张表存下来(table[i][ck])
2. 计算查询向量和样本向量的距离时，可以用距离聚类中心的距离来代替，进行O(1)的查表即可获取，（样本空间压缩后的编码是[c0, c1, c2... cM-1]，只要取之前计算好的table[i][ck]的距离即可）
3. 取最近K个即可

IVFPQ的工作原理

1
2
3
4
5
6
7
8
9
10
11
# 建立索引

1. 样本空间先用IVF 分类成nlist个聚类
2. 每个聚类中的样本向量计算它们和聚类中心的差值，得到新的向量
3. 使用PQ压缩每个聚类中新的样本向量

# 查询
1. 先找到距离查询量最近的聚类中心
2. 计算查询向量和聚类中心的差值，得到新的查询向量
3. 对新的查询向量使用PQ量化，并和该聚类中的压缩过后的样本向量做距离计算，并取出K个

代码示例

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import numpy as np
import faiss

# 生成一些示例数据
d = 128  # 向量维度
nb = 100000  # 数据库大小
nq = 10  # 查询数量
np.random.seed(1234)  # 确保可重复性
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')

# 构建IVFPQ索引
nlist = 100  # 聚类数量, ivf的聚类
m = 8  # 每个向量的子向量数量
nbit = 8 # ksub = 2^8, 每组要有ksub个聚类中心
quantizer = faiss.IndexFlatL2(d)  # 量化器
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, nbit)  # 8位量化

# 训练索引
index.train(xb)
index.add(xb)

# 查询
k = 4  # 返回的最近邻数量
D, I = index.search(xq, k)  # 搜索

# 打印结果
print("查询结果 (距离, 索引):")
for i in range(nq):
    print(f"查询 {i}:

总结

PQ量化压缩技术在向量索引领域较为常见，了解其整体的工作流程有助于我们选择合适的向量索引，从而达到较好的搜索性能. 而且PQ可以和很多算法结合，可以帮助理解部分faiss中index_factory的参数

▲