IVF_PQ (Inverted File Index with Product Quantization) 是一种结合了倒排文件索引 (IVF) 和乘积量化 (PQ) 的高级近似最近邻 (ANN) 算法。它旨在同时解决大规模高维向量数据的存储效率和查询速度问题，是目前许多主流向量数据库（如 Faiss、Milvus）中常用且高效的索引类型之一。IVF_PQ 通过“粗粒度 + 细粒度”的两阶段搜索策略，在保证较高召回率的同时，显著降低了内存占用和查询延迟。

核心思想：首先利用 IVF 的分区策略进行粗粒度筛选，将搜索范围缩小到少数几个簇；然后在这些选定的簇中，利用 PQ 对残差向量进行高效压缩和加速距离计算，实现高效的细粒度搜索。

一、为什么需要 IVF_PQ？结合两者的优势

我们已经了解了 IVF 和 PQ 各自的优势和局限：

  • IVF (倒排文件索引)

    • 优势：通过将数据划分为多个簇（Voronoi 单元）和构建倒排列表，实现了对搜索空间的粗粒度筛选，大大加速了查询速度。
    • 局限：如果倒排列表中存储的是原始浮点向量，随着向量维度和数据集规模的增大，内存占用仍然会非常高，且在倒排列表内部的精确距离计算依然是 $O((N/nlist) \cdot D)$，可能仍是瓶颈。

  • PQ (乘积量化)

    • 优势：通过向量分解和子向量独立量化，实现了对高维向量的极高压缩率，显著减少内存占用并加速距离计算（通过查表）。
    • 局限：直接使用 PQ 进行全局搜索时，由于其有损压缩特性，召回率可能不够理想，尤其是在整个空间进行近似距离计算时。

IVF_PQ 的出现正是为了结合两者的优势，弥补各自的不足：

  • IVF 负责将整个数据集划分，将相似的向量聚集在一起，减少了 PQ 需要处理的“噪声”和歧义，为 PQ 提供了一个更“干净”、更有局部特性的子集。
  • PQ 则负责在 IVF 选定的局部区域内，对向量进行高效压缩和距离计算，从而进一步降低内存和计算开销。

通过这种组合，IVF_PQ 能够在保证相对较高精度的前提下，实现对大规模高维向量数据的超高压缩率和极快查询速度。

二、IVF_PQ 的核心原理

IVF_PQ 的工作原理可以分解为索引构建和查询两个阶段，每个阶段都融合了 IVF 和 PQ 的思想。

2.1 索引构建阶段

  1. 全局聚类 (Coarse Quantization)：

    • 首先，对整个数据集（或其代表性子集）运行 K-Means 算法，得到 nlist 个簇中心 (centroids)。这些簇中心将整个向量空间划分为 nlist 个 Voronoi 单元。这一步与纯 IVF 索引的构建相同。
    • Faiss 中，这个阶段被称为 Quantizer 的训练。

  2. 残差向量计算 (Residual Vector Computation)：

    • 对于数据库中的每个原始向量 $\mathbf{x}$，首先找到其最近的簇中心 $\mathbf{c}_j$。
    • 然后，计算该向量相对于其簇中心的残差向量 (residual vector)：$\mathbf{r} = \mathbf{x} - \mathbf{c}_j$。
    • 残差向量的目的是捕获向量在所属簇内的局部信息，因为簇中心已经捕获了大部分全局信息。

  3. 残差向量的乘积量化 (Product Quantization of Residuals)：

    • 对于所有残差向量组成的集合，进行全局的乘积量化训练。这意味着对所有残差向量，将其分解为 $M$ 个子向量，并为每个子向量空间独立地训练一个子码本。
    • 训练完成后，每个子码本包含 $K$ 个码字。
    • 重要：这里的 PQ 训练是在残差向量空间上进行的，而不是原始向量空间。

  4. 倒排列表构建与编码 (Inverted List Construction and Encoding)：

    • 构建 nlist 个倒排列表。
    • 对于每个原始向量 $\mathbf{x}$，在找到其最近的簇中心 $\mathbf{c}_j$ 后，计算其残差向量 $\mathbf{r} = \mathbf{x} - \mathbf{c}_j$。
    • 将残差向量 $\mathbf{r}$ 使用训练好的 PQ 编码器进行编码，得到一个由 $M$ 个码字索引组成的短序列：$\text{code}(\mathbf{r}) = [idx_1, idx_2, \ldots, idx_M]$。
    • 将这个编码后的短序列 $\text{code}(\mathbf{r})$（通常是 $M$ 字节）存储到簇中心 $\mathbf{c}_j$ 对应的倒排列表中。

2.2 查询阶段

IVF_PQ 的查询过程是一个两阶段的近似搜索：

  1. 粗粒度搜索 (Coarse Search / Quantizer Search)：

    • 给定一个查询向量 $Q$，首先计算 $Q$ 到所有 nlist 个簇中心的距离。
    • 选择距离 $Q$ 最近的 nprobe 个簇中心，这些中心对应的倒排列表将构成细粒度搜索的候选空间。

  2. 细粒度搜索 (Fine Search / Product Quantization Search)：

    • 对于每个被选中的 nprobe 个簇中心 $\mathbf{c}_j$：
      • 计算查询向量 $Q$ 相对于该簇中心的残差查询向量：$Q_r = Q - \mathbf{c}_j$。
      • 使用非对称距离计算 (ADC) 策略：预计算 $Q_r$ 的 $M$ 个子向量到 PQ 各个子码本中所有码字的距离，生成一个距离查找表 (DLT)。
      • 遍历该簇中心 $\mathbf{c}_j$ 对应的倒排列表。对于列表中的每个编码残差向量 $\text{code}(\mathbf{r}) = [idx_1, idx_2, \ldots, idx_M]$：
        • 通过查表，快速计算 $Q_r$ 与 $\text{code}(\mathbf{r})$ 对应的近似距离：$D(Q, \mathbf{x}) \approx \sum_{m=1}^{M} D(C_m[idx_m], Q_{r,m})$。
    • 将所有从 nprobe 个倒排列表中获得的近似距离和对应的向量 ID 进行聚合。

  3. 结果聚合与排序 (Result Aggregation and Sorting)：

    • 从所有候选结果中，选择距离查询向量 $Q$ 最近的 Top-K 个向量作为最终结果。

IVF_PQ 索引概念图

说明：IVF_PQ 构建时，首先通过 K-Means 划分簇，然后对所有向量的残差进行 PQ 训练。每个向量存储时，不再是原始向量或残差向量，而是其残差向量的 PQ 编码。查询时，先通过查询向量与簇中心的距离确定搜索范围 (nprobe 个簇)，然后在这些簇内，通过预计算的距离查找表 (DLT) 快速计算查询残差与存储的 PQ 编码残差之间的近似距离。

三、IVF_PQ 的参数调优

IVF_PQ 的性能和精度受多个参数影响，主要包括：

  1. nlist (Number of Centroids / Clusters)：

    • 与纯 IVF 相同，决定了粗粒度分区的数量。
    • 影响：更大的 nlist 意味着更细的分区，每个倒排列表更短，通常有助于 PQ 在局部区域更好地工作，可能提高精度，但索引构建更慢，且查找 nprobe 个簇中心时的开销可能更大。
    • 选择建议：通常 nlist 介于 $\sqrt{N}$ 到 $N/39$ 之间，具体取决于数据集大小和分布。

  2. M (Number of subvectors for PQ)：

    • 与纯 PQ 相同，决定了残差向量被分解成子向量的数量。
    • 影响：
      • 更大的 M：压缩率更高，每个编码残差向量占用字节数更少，查询时的查表次数增多。子向量维度 $D/M$ 降低，可能影响 PQ 训练精度。
      • 更小的 M：压缩率更低，每个编码残差向量占用字节数更多，子向量维度 $D/M$ 升高，PQ 训练精度可能更高。
    • 选择建议：通常 $M$ 的值在 4 到 64 之间，使得 $D/M \ge 16$ 左右。常见的 $M=8$ 或 $M=16$。

  3. nbits (Number of bits per subvector for PQ)：

    • 与纯 PQ 相同，决定了每个子码本中的码字数量 $K = 2^{nbits}$。
    • 影响：几乎总是取 nbits=8 (即 $K=256$)，因为 1 字节索引是最高效的存储方式，且能提供合理的精度。
    • 选择建议：通常固定为 8。

  4. nprobe (Number of Probed Inverted Lists)：

    • 与纯 IVF 相同，决定了查询时要检查的倒排列表的数量。
    • 影响：
      • 更大的 nprobe：召回率（精度）更高，但查询时间更长。
      • 更小的 nprobe：查询速度更快，但召回率可能下降。
    • 选择建议：在运行时动态调整以平衡精度和速度的关键参数。

四、IVF_PQ 的优势与局限性

4.1 优势

  1. 高压缩率：通过 PQ 对残差向量进行编码，极大减少了每个向量的存储空间，从而允许在内存中加载数十亿级的向量数据。
  2. 超快查询速度：结合了 IVF 的粗粒度筛选和 PQ 的细粒度查表距离计算，查询时间通常是近似对数级。
  3. 良好的精度-速度-内存平衡：通过细致的参数调优（特别是 nlist 和 nprobe），可以在极低的内存占用和极高的查询速度下，实现可接受甚至优秀的召回率。
  4. 适用于超大规模数据集：是处理 TB 级甚至 PB 级高维向量数据的首选索引类型之一。

4.2 局限性

  1. 索引构建复杂且耗时：需要先进行 K-Means 聚类训练 (IVF 部分)，再对所有残差向量进行 PQ 训练。这两个训练阶段都可能非常耗时，尤其是在数据集非常大时。
  2. 有损压缩，存在精度损失：残差向量的 PQ 编码是有损的，必然会引入量化误差，从而导致查询结果并非 100% 精确。
  3. 参数调优复杂：nlist, nprobe, M, nbits 等参数的组合和选择对最终性能和精度影响巨大，需要经验和实验来找到最优配置。
  4. 动态更新成本高：与纯 IVF 类似，增加或删除向量可能需要重新计算残差并重新编码，甚至在大量更新后需要重新训练聚类中心，影响了索引的动态性。

五、Python 示例：使用 Faiss 实现 IVF_PQ 索引

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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import faiss
import numpy as np
import time

# 1. 生成模拟数据
D = 128      # 向量维度
NB = 1000000 # 数据库中的向量数量 (100万)
NQ = 10      # 查询向量数量

# 随机生成数据库向量 (标准化到单位长度，以便L2距离能近似表示余弦相似度)
xb = np.random.random((NB, D)).astype('float32')
xb = xb / np.linalg.norm(xb, axis=1)[:, np.newaxis]

# 随机生成查询向量
xq = np.random.random((NQ, D)).astype('float32')
xq = xq / np.linalg.norm(xq, axis=1)[:, np.newaxis]

print(f"数据库向量形状: {xb.shape}")
print(f"查询向量形状: {xq.shape}")

# --- 2. 创建 IVF_PQ 索引 ---
# nlist: 簇的数量 (IVF部分)
nlist = 1000 
# M: 子向量的数量 (PQ部分，这里将128维向量分成8个16维子向量)
M = 8 
# nbits: 每个子向量的量化比特数 (PQ部分，2^8 = 256 个码字)
nbits = 8 

# Quantizer (量化器): 用于对向量进行粗粒度聚类，找到簇中心。
# IndexFlatL2 意味着量化器使用精确的L2距离查找最近的簇中心。
quantizer = faiss.IndexFlatL2(D) 

# IndexIVFPQ 构造函数:
# - quantizer: 量化器 (IndexFlatL2)
# - D: 向量维度
# - nlist: 簇的数量 (IVF)
# - M: 子向量数量 (PQ)
# - nbits: 每个子向量的量化比特数 (PQ)
# - metric: 距离度量 (PQ通常用于L2或内积)
index_ivfpq = faiss.IndexIVFPQ(quantizer, D, nlist, M, nbits, faiss.METRIC_L2)

# --- 3. 训练 IVF_PQ 索引 ---
# IVF_PQ 索引需要进行两次训练：
# 1. Quantizer (IVF部分) 的训练：确定 nlist 个簇中心
# 2. PQ (PQ部分) 的训练：为残差向量训练 M 个子码本
print(f"\n开始训练 IVF_PQ 索引 (nlist={nlist}, M={M}, nbits={nbits})...")
start_time = time.time()
index_ivfpq.train(xb) # 训练数据通常是整个数据集或其一个代表性子集。
end_time = time.time()
print(f"IVF_PQ 索引训练耗时: {end_time - start_time:.4f} 秒")

# 检查索引是否已经训练好
print(f"索引是否已训练: {index_ivfpq.is_trained}")

# --- 4. 添加向量到 IVF_PQ 索引 ---
# 在 add 过程中，每个向量首先被分配到最近的簇，然后其残差向量被PQ编码并存储。
print("添加向量到 IVF_PQ 索引 (编码过程)...")
start_time = time.time()
index_ivfpq.add(xb)
end_time = time.time()
print(f"IVF_PQ 索引添加向量耗时: {end_time - start_time:.4f} 秒")
print(f"索引中向量数量: {index_ivfpq.ntotal}")

# 计算索引的内存占用（近似）
# 每个向量占用 M 字节 (存储PQ编码的残差) + Quantizer的开销
estimated_memory_bytes = index_ivfpq.ntotal * M + nlist * D * 4 # M字节编码 + nlist个D维中心
print(f"IVF_PQ 索引估计内存占用: {estimated_memory_bytes / (1024**2):.2f} MB")
# 原始向量内存占用 (100万 * 128 * 4 字节)
original_memory_bytes = NB * D * 4
print(f"原始向量内存占用: {original_memory_bytes / (1024**2):.2f} MB")
print(f"压缩比: {original_memory_bytes / (index_ivfpq.ntotal * M):.2f}X (仅考虑数据压缩)") # 更准确的压缩比应考虑整个索引结构

# --- 5. 查询 IVF_PQ 索引 ---
k = 5 # 查找最近的5个邻居

# 设置查询参数 nprobe (查询时检查的簇的数量)
# nprobe 越大，精度越高，但查询越慢。
index_ivfpq.nprobe = 10 # 初始设置检查10个最近的簇

print(f"\n开始查询 (nprobe={index_ivfpq.nprobe})...")
start_time = time.time()
distances_ivfpq, indices_ivfpq = index_ivfpq.search(xq, k)
end_time = time.time()

print(f"IVF_PQ 近似搜索耗时 (nprobe={index_ivfpq.nprobe}): {end_time - start_time:.4f} 秒")
print("查询向量 0 的近似最近邻索引:", indices_ivfpq[0])
print("查询向量 0 的近似最近邻距离:", distances_ivfpq[0])


# --- 6. 评估召回率 (与精确搜索比较) ---
# 为了评估召回率，我们首先进行一个精确搜索作为基准
index_exact = faiss.IndexFlatL2(D)
index_exact.add(xb)
_, indices_exact = index_exact.search(xq, k)

correct_matches_initial = np.isin(indices_ivfpq[0], indices_exact[0]).sum()
recall_initial = correct_matches_initial / k
print(f"\n查询向量 0 的召回率 (nprobe={index_ivfpq.nprobe}): {recall_initial:.4f}")

# 提高 nprobe，观察对速度和召回率的影响
index_ivfpq.nprobe = 50 

print(f"\n再次查询 (nprobe={index_ivfpq.nprobe}, 提高精度)...")
start_time = time.time()
distances_ivfpq_high_nprobe, indices_ivfpq_high_nprobe = index_ivfpq.search(xq, k)
end_time = time.time()
print(f"近似搜索耗时 (nprobe={index_ivfpq.nprobe}): {end_time - start_time:.4f} 秒")

correct_matches_high_nprobe = np.isin(indices_ivfpq_high_nprobe[0], indices_exact[0]).sum()
recall_high_nprobe = correct_matches_high_nprobe / k
print(f"查询向量 0 的召回率 (nprobe={index_ivfpq.nprobe}): {recall_high_nprobe:.4f}")

示例解读：

  • IndexIVFPQ(quantizer, D, nlist, M, nbits, metric)：创建 IVF_PQ 索引。这里的参数融合了 IVF 和 PQ 的特性。
  • 训练 (train)：IVF_PQ 的训练过程比纯 IVF 或纯 PQ 更复杂，它需要同时完成粗粒度聚类（quantizer 部分）和残差向量的乘积量化训练（PQ 部分）。
  • 添加 (add)：在添加向量时，每个向量首先被归属到最近的簇，然后计算残差，残差再被 PQ 编码成 M 字节的短序列，存储在对应的倒排列表中。
  • 内存占用：可以看到，IVF_PQ 的内存占用比原始向量显著减少，是其主要优势之一。
  • nprobe 调优：与纯 IVF 类似，nprobe 是在查询时动态调整精度和速度的关键参数。通过增加 nprobe，可以提高召回率，但查询时间也会增加。

六、总结

IVF_PQ 是目前最强大、最通用的近似最近邻算法之一，广泛应用于需要处理大规模、高维向量数据的场景。它通过巧妙地结合 IVF 的分区策略和 PQ 的向量压缩及距离计算加速，实现了存储效率、查询速度和召回率之间的卓越平衡。理解 IVF_PQ 的工作原理、参数调优及其优势与局限性，是构建高性能向量搜索系统和处理海量非结构化数据挑战的关键。