Faiss 与 PQ:压缩的艺术与科学
你好!如果你正在和海量的向量数据打交道,并且想用 Faiss 来加速你的相似性搜索,那你一定听说过或者正在使用 PQ(Product Quantization,乘积量化)。这玩意儿简直是处理大规模向量数据的瑞士军刀——既能大幅压缩向量占用的内存,又能显著提升搜索速度。但是,就像任何强大的工具一样,用好 PQ 需要一点技巧和理解。
很多人可能只是简单地用了 IndexPQ
或者 IndexIVFPQ
,效果好像还行,但总觉得没发挥出它的全部潜力。或者更糟,参数没设对,导致精度暴跌,或者内存/速度优化效果不明显。别担心,今天我们就来深入聊聊 Faiss 里的 PQ,特别是那些关键参数到底是怎么回事,以及如何在不同场景下把它们调到最佳状态。
咱们的目标是让你不仅知道 PQ 是什么,更能掌握如何在 Faiss 里把它玩得明明白白,成为真正的“Faiss 大法师”!
PQ 到底是个啥?压缩的魔法
想象一下,你有一堆高维向量,比如 128 维、256 维甚至更高。直接存储和计算它们的距离,对于上亿甚至几十亿的向量来说,内存和计算成本都高得吓人。怎么办?压缩!
PQ 的想法非常巧妙。它不是对整个向量进行压缩,而是先把高维向量“切”成几段(子向量),然后对每一段分别进行量化。
- 向量切分 (Splitting): 假设你有一个 D 维向量
x
。PQ 会把它切成M
个 D/M 维的子向量。比如,一个 128 维向量,如果M=8
,就会被切成 8 个 16 维的子向量。 - 子空间量化 (Quantization in Subspaces): 接下来,对每个子空间(也就是所有向量的第
i
段子向量构成的集合)单独进行 K-Means 聚类。通常,我们会为每个子空间找到k* = 2^nbits
个聚类中心(centroids)。nbits
通常是 8,意味着每个子空间有 256 个聚类中心。这些聚类中心就构成了这个子空间的“码本”(codebook)。 - 编码 (Encoding): 对于原始向量
x
的每个子向量x_i
,找到它在对应子空间码本中最接近的那个聚类中心的 ID (0 到 k*-1)。最后,原始向量x
就被压缩成了一个由M
个小整数(码字 ID)组成的列表。比如,如果M=8
,nbits=8
,那么一个 128 维的浮点向量就被压缩成了 8 个 8-bit 整数,总共只需要 8 字节!相比原始的 128 * 4 = 512 字节(假设是 float32),压缩率惊人。
思考一下: 为什么不直接对整个 128 维向量做 K-Means 呢?如果你想用 8 字节(64 bit)来编码,意味着你需要 2^64
个聚类中心!这在计算和存储上都是完全不可行的。PQ 通过子空间分解,用 M * k*
个聚类中心(例如 8 * 256 = 2048 个)就巧妙地实现了高压缩率。这就是它的魔力所在:用可控的码本大小,实现了指数级的编码空间。
Faiss 中的 PQ 实现:IndexPQ
与 IndexIVFPQ
Faiss 提供了几种利用 PQ 的索引类型。最基础的是 IndexPQ
。
import faiss
import numpy as np
d = 128 # 向量维度
nb = 100000 # 数据库向量数量
nq = 10000 # 查询向量数量
# 生成一些随机数据
np.random.seed(1234)
x_train = np.random.random((nb // 10, d)).astype('float32') # 训练数据
x_base = np.random.random((nb, d)).astype('float32') # 数据库向量
x_query = np.random.random((nq, d)).astype('float32') # 查询向量
M = 8 # 子向量数量
nbits = 8 # 每个子向量索引的比特数
# 创建 IndexPQ
pq_index = faiss.IndexPQ(d, M, nbits)
# 训练码本 (非常重要!)
print("Training PQ...")
pq_index.train(x_train) # 使用训练数据来学习子空间的聚类中心
# 添加数据库向量 (Faiss 会自动进行编码)
print("Adding vectors...")
pq_index.add(x_base)
print(f"Index total vectors: {pq_index.ntotal}")
print(f"Is trained: {pq_index.is_trained}")
print(f"Code size (bytes per vector): {pq_index.code_size}") # 这就是 M * nbits / 8
# 搜索
k = 10 # 查找最近的 10 个邻居
print("Searching...")
D, I = pq_index.search(x_query, k)
# D 是距离,I 是索引
# print(I[:5])
# print(D[:5])
IndexPQ
是最纯粹的 PQ 实现。它直接计算查询向量和数据库中所有 PQ 编码向量之间的近似距离,然后返回距离最小的 k 个。这种距离计算通常是非对称距离计算 (Asymmetric Distance Computation, ADC):
- 查询向量
q
不进行压缩。 - 数据库向量
x
是 PQ 编码后的code(x)
。 - 计算
d(q, x)
的近似值d'(q, code(x))
。
具体怎么算呢?对于查询向量 q
,也把它切成 M
个子向量 q_i
。然后,对于每个子向量 q_i
,预先计算它到对应子空间码本中所有 k*
个聚类中心 c_{ij}
的距离 d(q_i, c_{ij})
。这个距离表可以预计算并缓存。
当计算 q
和某个编码为 (id_1, id_2, ..., id_M)
的数据库向量 x
的距离时,只需要查表并将各子空间的距离加起来:
d'(q, code(x)) = Σ_{i=1 to M} d(q_i, c_{i, id_i})
这种查表相加的方式非常快!这就是 PQ 加速搜索的核心原理之一。
但是,IndexPQ
有个巨大的缺点: 它需要计算查询向量与数据库中 所有 向量的距离。当数据库非常大时(比如上亿),即使 ADC 计算很快,这个线性扫描的开销依然是无法接受的。
所以,实践中更常用的是 IndexIVFPQ
。它结合了 IVF(Inverted File,倒排文件)和 PQ。
- IVF 的作用: 先用 K-Means 将整个向量空间粗略地划分成
nlist
个单元(Voronoi cells)。每个数据库向量被分配到离它最近的那个单元中心对应的倒排列表中。搜索时,查询向量q
先找到离它最近的nprobe
个单元,然后只需要在这些单元对应的倒排列表里进行搜索,而不需要扫描整个数据库。 - PQ 的作用: 在每个倒排列表内部,存储的不是原始向量,而是它们的 PQ 编码。这样既节省了存储空间,又可以在选定的
nprobe
个列表内进行快速的 ADC 距离计算。
import faiss
import numpy as np
d = 128
nb = 100000
nq = 10000
np.random.seed(1234)
x_train = np.random.random((nb // 10, d)).astype('float32')
x_base = np.random.random((nb, d)).astype('float32')
x_query = np.random.random((nq, d)).astype('float32')
M = 8
nbits = 8
nlist = 100 # IVF 的单元数量
quantizer = faiss.IndexFlatL2(d) # 用于 IVF 划分的粗量化器,这里用精确 L2 距离
# 创建 IndexIVFPQ
ivfpq_index = faiss.IndexIVFPQ(quantizer, d, nlist, M, nbits)
# 训练 (需要同时训练 IVF 的聚类中心和 PQ 的码本)
print("Training IVFPQ...")
ivfpq_index.train(x_train)
# 添加
print("Adding vectors...")
ivfpq_index.add(x_base)
print(f"Index total vectors: {ivfpq_index.ntotal}")
print(f"Is trained: {ivfpq_index.is_trained}")
# 设置搜索时探索的单元数量 (关键参数!)
ivfpq_index.nprobe = 10
# 搜索
k = 10
print("Searching...")
D, I = ivfpq_index.search(x_query, k)
# print(I[:5])
# print(D[:5])
IndexIVFPQ
是 Faiss 中最常用的大规模向量索引之一。它的性能受到 IVF 参数(nlist
, nprobe
)和 PQ 参数(M
, nbits
)的共同影响。
核心参数调优:M 和 nbits
现在进入正题,我们来仔细看看 PQ 的两个核心参数:M
和 nbits
。
M
(Number of subquantizers): 子向量的数量。它直接决定了向量被切成多少段。d
必须能被M
整除。M
越大,每个子向量的维度dsub = d / M
就越小。nbits
(Number of bits per subquantizer index): 每个子向量用多少比特来编码其最近的码本中心 ID。通常取 8,意味着每个子空间有k* = 2^8 = 256
个中心。也可以取其他值,比如 4 到 16 之间,但 8 是最常见的选择。
这两个参数共同决定了 PQ 的压缩率和精度,以及对内存和速度的影响。
1. 内存占用 (Code Size):
每个向量的 PQ 编码大小是 M * nbits / 8
字节。这是最直观的影响。
- 增大
M
或nbits
都会增加内存占用。 - 对于固定维度
d
,通常我们希望在满足精度要求的前提下,尽量减小M * nbits
。
码本内存占用: 别忘了还有码本本身也需要存储!码本的大小是 M * k* * dsub * sizeof(float)
字节,即 M * (2^nbits) * (d / M) * 4 = d * 2^nbits * 4
字节 (假设是 float32)。
- 码本大小与
M
无关!它主要由d
和nbits
决定。 - 如果
nbits
很大(比如 12 或 16),码本可能会占用相当大的内存,尤其是在分布式 Faiss 场景下,每个节点可能都需要一份码本的拷贝。
2. 精度 (Accuracy):
PQ 本质上是一种有损压缩,精度损失是必然的。M
和 nbits
如何影响精度呢?
- 增大
M
: 每个子向量的维度dsub = d / M
变小。在维度更低的空间里做 K-Means,通常更容易找到好的聚类中心来表示数据,从而减少量化误差。因此,增大M
通常会提高精度。 - 增大
nbits
: 每个子空间有更多的聚类中心 (k* = 2^nbits
)。更多的中心意味着子向量可以被更精确地表示,量化误差减小。因此,增大nbits
通常也会提高精度。
思考与权衡: 提高精度通常意味着牺牲压缩率(更大的 M * nbits
)和/或增加码本大小(更大的 nbits
)。这是一个典型的精度-资源权衡。
3. 搜索速度 (ADC Speed):
ADC 的速度主要取决于查表相加操作的次数,即 M
次查表和 M-1
次加法。因此,M
越大,ADC 计算理论上越慢。但实际影响也跟 CPU 缓存、内存带宽等因素有关。
nbits
对 ADC 速度的影响主要体现在预计算距离表 d(q_i, c_{ij})
的大小和计算时间。这个表的大小是 M * k* = M * 2^nbits
。如果这个表能完全放入 CPU 缓存(比如 L1/L2 cache),那么查表会非常快。如果 M
或 nbits
增大导致表变得太大,缓存命中率下降,速度可能会受到影响。
常见误区: 有人认为 M
越大精度越高,所以应该越大越好。但别忘了它会增加编码大小和 ADC 计算时间。实践中 M
的选择需要在精度和资源消耗之间找到平衡点。
参数选择的经验法则和场景分析:
nbits
的选择:nbits = 8
是最常用、通常也是效果不错的起点。它提供了 256 个中心,对于多数子空间维度来说足够了。码本大小d * 256 * 4
字节,相对可控。- 如果内存极其有限,可以尝试
nbits = 4
或6
。但这通常会导致精度显著下降,因为每个子空间只有 16 或 64 个中心。 - 如果对精度要求非常高,并且能接受更大的码本内存占用,可以尝试
nbits = 10
或12
(1024 或 4096 个中心)。超过 12 通常带来的精度提升有限,但码本内存和预计算开销会急剧增加。 nbits > 8
在IndexPQ
中可能效果更明显,但在IndexIVFPQ
中,由于 IVF 已经筛选了一部分候选向量,PQ 部分的精度要求可能没那么极端,nbits = 8
往往足够。
M
的选择:M
的选择与原始维度d
密切相关。目标是让子向量维度dsub = d / M
不至于太高(比如 > 20 可能就偏高了),也不至于太低(比如 < 4 可能信息损失过大)。- 常见选择:
- 对于
d = 128
,M
可以选 8 (dsub=16), 16 (dsub=8), 32 (dsub=4)。 - 对于
d = 256
,M
可以选 16 (dsub=16), 32 (dsub=8), 64 (dsub=4)。 - 对于
d = 512
,M
可以选 32 (dsub=16), 64 (dsub=8), 128 (dsub=4)。
- 对于
- 权衡:
- 高精度场景: 优先考虑增大
M
(比如dsub
降到 4 或 8),同时保持nbits=8
。例如,d=128, M=16
或M=32
。这会增加编码大小,但精度通常更好。 - 内存/速度优先场景: 优先考虑减小
M
(比如dsub
升到 16),同时保持nbits=8
。例如,d=128, M=8
。或者在M
固定的情况下,尝试降低nbits
(如nbits=6
),但这通常精度损失更大。 - 一个经验: 对于给定的
d
,M
值使得dsub
在 8 到 16 之间通常是一个不错的起点。然后根据具体需求调整M
。
- 高精度场景: 优先考虑增大
- 流式思考: 我拿到一个 192 维的向量,内存预算比较紧。
d=192
。M
必须能整除 192。可以选M=12
(dsub=16),M=16
(dsub=12),M=24
(dsub=8),M=32
(dsub=6),M=48
(dsub=4)。如果内存优先,M=12
(12字节/向量,如果 nbits=8) 是个选择,dsub=16 精度可能一般。M=16
(16字节/向量) 看起来不错,dsub=12。M=24
(24字节/向量) 精度可能更好,dsub=8。如果 16 字节可以接受,我会先尝试M=16, nbits=8
。如果不行,再看是牺牲更多内存换精度 (M=24
) 还是牺牲精度换内存 (M=12
或M=16, nbits=6
)。
训练数据量: PQ 的训练(学习码本)也需要数据。Faiss 官方建议每个子空间的 K-Means 至少需要
k* * 39
个训练样本 (来自 K-Means 的经验规则)。所以,总训练样本数至少需要M * k* * 39
?不对! 训练是并行地在M
个子空间上进行的。你只需要保证总训练样本量N_train
足够大,使得每个子空间分配到的样本N_train
足够进行 K-Means 训练。通常,几十万甚至几百万的训练向量是比较安全的,特别是当nbits
较大时。
IndexIVFPQ
中的额外考量:与 IVF 参数的互动
当你使用 IndexIVFPQ
时,PQ 参数还需要和 IVF 参数 nlist
、nprobe
结合起来看。
nlist
(Number of Voronoi cells): 倒排列表的数量。通常建议取数据库规模nb
的平方根附近,比如4*sqrt(nb)
到16*sqrt(nb)
之间。nlist
太小,每个列表包含太多向量,IVF 的过滤效果不明显;nlist
太大,可能很多列表为空或只有少量向量,浪费内存,且nprobe
需要设置得更大才能保证召回率。nprobe
(Number of probes): 搜索时检查的单元数量。这是速度 vs 精度的关键权衡点!nprobe
越大,搜索范围越大,召回率(精度)越高,但速度越慢。nprobe
越小,速度越快,但可能错过正确的近邻,导致召回率下降。
PQ 参数与 nprobe
的关系:
想象一下,IVF 负责“粗筛”,PQ 负责在粗筛后的列表里“精排”(虽然是近似距离)。
- 如果你的 PQ 参数设置得比较“激进”(比如
M
较小,nbits
较小,压缩率高但精度损失大),那么即使 IVF 找到了包含真正近邻的单元,PQ 的近似距离计算也可能出错,导致排序错误。这种情况下,你可能需要增大nprobe
来弥补 PQ 带来的精度损失,通过检查更多的单元来增加找到真正近邻的机会。 - 反之,如果你的 PQ 参数设置得比较“保守”(比如
M
较大,nbits=8
或更高,精度较高),那么 IVF 筛选出的单元内的 PQ 近似距离排序会更可靠。这种情况下,你可以用较小的nprobe
来获得更快的速度,同时仍然保持不错的召回率。
调优策略 IndexIVFPQ
:
- 确定内存预算: 这通常会限制
M * nbits / 8
的大小。 - 选择
nbits
: 从nbits=8
开始。 - 选择
M
: 根据维度d
和内存预算选择M
,使dsub = d / M
在一个合理范围 (e.g., 4-16)。 - 选择
nlist
: 根据数据库大小nb
估算一个nlist
值 (e.g.,k * sqrt(nb)
,k 在 4 到 16 之间)。 - 训练索引:
index.train(x_train)
- 添加数据:
index.add(x_base)
- 调优
nprobe
: 这是最关键的一步!通过在一个标注了真实近邻的验证集上进行测试,绘制 Recall@k vs. Query Time (或nprobe
) 的曲线。选择一个能在可接受的查询时间内达到你的目标召回率的nprobe
值。 - 迭代优化: 如果在可接受的
nprobe
范围内,召回率始终达不到要求,可以考虑:- 增加 PQ 精度:增大
M
(如果内存允许) 或尝试nbits=10/12
(如果码本内存允许)。 - 调整
nlist
:如果nprobe
已经很大但召回率仍不高,可能nlist
太大了,试试减小nlist
;如果nprobe=1
时速度已经很慢,可能nlist
太小了,试试增大nlist
。 - 重新训练:确保训练数据足够且有代表性。
- 增加 PQ 精度:增大
一个真实的例子: 假设你有 1 亿 128 维向量 (nb=1e8, d=128
),目标是 Recall@10 > 90%,查询时间 < 10ms。
- 内存预算:假设每个向量最多 16 字节。
M * nbits / 8 <= 16
。 - 尝试
nbits=8
。则M <= 16
。 - 选择
M=16
(dsub=8)。编码大小 16 字节。 - 选择
nlist
。sqrt(1e8) = 10000
。可选nlist
在 40000 到 160000 之间。假设选nlist = 65536
(2^16
)。 - 训练,添加数据。
- 测试不同
nprobe
值:nprobe=1
: Recall@10 = 60%, Time = 1msnprobe=5
: Recall@10 = 80%, Time = 4msnprobe=10
: Recall@10 = 88%, Time = 7msnprobe=20
: Recall@10 = 92%, Time = 13ms
- 分析:
nprobe=20
满足了召回率要求,但超时了。nprobe=10
接近目标。 - 优化方向:
- 方案 A (牺牲一点内存换精度): 保持
nlist=65536
,尝试M=32
(dsub=4, 编码 32 字节)。期望 PQ 精度提高,可能用nprobe=10
就能达到 90% 召回率,且时间可能仍在 10ms 左右 (因为M
增大了 ADC 变慢,但nprobe
不变)。需要重新测试。 - 方案 B (调整 IVF): 保持
M=16, nbits=8
。尝试减小nlist
,比如nlist=32768
。每个列表更大,可能用更小的nprobe
(比如nprobe=15
) 就能达到 90% 召回率,且总时间可能 < 10ms。需要重新训练和测试。 - 方案 C (终极手段): 如果上述都不行,可能需要牺牲更多内存(更大的 M)或接受更低的召回率/更长的查询时间,或者考虑更复杂的索引结构 (如 HNSW + PQ)。
- 方案 A (牺牲一点内存换精度): 保持
其他技巧和注意事项
- 对称距离 (SDC) vs 非对称距离 (ADC): Faiss 的 PQ 默认使用 ADC (查询向量不压缩)。也可以配置成 SDC (查询向量也进行 PQ 编码),但通常 ADC 效果更好。
- PQ 训练的稳定性: K-Means 对初始化敏感。Faiss 内部有多次运行 K-Means (nredo) 并选择最佳结果的机制,有助于提高码本质量。
- 预计算距离表: Faiss 会自动处理 ADC 中距离表的预计算和缓存。理解其原理有助于分析性能瓶颈。
- Refinement (重排): 对于
IndexIVFPQ
,可以在得到 PQ 近似距离排序的 Top N 个结果后,再用这些向量的原始向量(如果存储了的话,或者需要从外部存储读取)与查询向量计算精确距离,进行重新排序。这可以显著提高最终结果的精度,但会增加额外的计算和 IO 开销。 - 多 GPU/分布式: Faiss 支持在多 GPU 或多台机器上运行,PQ 码本如何在这些设备间同步和使用需要注意。
总结
Faiss 中的 PQ 是一种强大的向量压缩和加速技术,但用好它需要理解其核心参数 M
和 nbits
的影响以及它们之间的权衡。
M
(子向量数): 主要影响精度和 ADC 速度。M
越大,精度通常越高,但编码变大,ADC 变慢。nbits
(比特数/子向量): 主要影响精度和码本大小。nbits
越大,精度越高,码本越大,编码也变大。nbits=8
是常用起点。
在 IndexIVFPQ
中,PQ 参数需要与 IVF 参数 nlist
和 nprobe
协同调优。
- 根据内存预算和维度
d
初步选择M
和nbits
。 - 根据数据库大小
nb
选择nlist
。 - 通过实验确定最佳的
nprobe
,在速度和召回率之间找到平衡。 - 如果效果不达标,迭代调整
M
,nbits
,nlist
。
调优过程就像一场实验,需要耐心和对数据/场景的理解。没有万能的参数组合,只有最适合你特定需求的那一个。希望这篇深入探讨能帮助你更好地驾驭 Faiss PQ,让你的向量搜索更快、更准、更省资源!开始你的调优之旅吧!