你好！如果你正在使用 Faiss 处理大规模向量相似性搜索，或者对向量量化技术充满好奇，那么你一定听说过 Product Quantization (PQ，乘积量化)。PQ 是 Faiss 中一种极其重要的向量压缩和近似搜索技术。它如何在大幅降低内存占用的同时，还能提供不错的搜索精度？这背后的数学原理是什么？特别是，它是如何近似计算向量距离的？
今天，我们就来一场硬核的数学之旅，深入剖析 PQ 算法的原理，重点推导其核心的距离计算方法。准备好你的纸和笔，让我们一起揭开 PQ 的神秘面纱！
为什么需要 PQ？ANN 的挑战
在探讨 PQ 之前，我们先简单回顾一下它要解决的问题。在高维空间中进行精确的最近邻搜索（Nearest Neighbor Search）通常计算量巨大，尤其当数据集规模达到百万、千万甚至数十亿级别时。想象一下，对于一个查询向量 q，你需要计算它与数据库中每一个向量 y 的距离，然后找到最小的那个。如果向量维度 D 很高，数据集大小 N 很大，这个 O(N*D) 的计算复杂度是难以接受的。
因此，近似最近邻搜索（Approximate Nearest Neighbor Search, ANN）应运而生。ANN 的目标是在牺牲一定精度的情况下，大幅提升搜索速度和降低内存占用。Faiss 提供了多种 ANN 算法，而 PQ 正是其中降低内存占用和加速距离计算的关键技术之一。
PQ 的核心思想是有损压缩。它将高维向量压缩成非常紧凑的编码（通常只需要几十个比特），并基于这些编码来近似计算原始向量间的距离。
PQ 的核心原理：分解与量化
PQ 的巧妙之处在于“乘积”（Product）这个词。它并不是直接对整个高维向量进行量化，而是先将向量分解成多个低维子向量，然后对每个子向量独立地进行量化。
假设我们有一个 D 维的向量 x。PQ 首先会将这个向量切分成 M 段（或者叫子空间），每段包含 D_sub = D / M 个维度。（这里假设 D 可以被 M 整除，如果不可以，通常会做一些填充处理）。
x = [x_1 | x_2 | ... | x_M]
其中 x_m 是一个 D_sub 维的子向量，m = 1, ..., M。
接下来，对每一个子空间 m，我们独立地运行一个聚类算法（通常是 K-means）来学习一个包含 k* 个聚类中心（centroids）的码本（codebook）。k* 通常选择为 256（即 2^8），这样每个子向量的量化编码就可以用 8 个比特（1 个字节）来表示。
对于第 m 个子空间，我们得到一个码本 C_m = {c_{m,1}, c_{m,2}, ..., c_{m,k*}}，其中 c_{m,j} 是一个 D_sub 维的质心向量。
训练过程：
- 从训练数据集中抽取大量样本向量。
- 将每个样本向量
x切分成M个子向量x_1, ..., x_M。 - 对于每个子空间
m(从 1 到M)：- 收集所有样本的第
m个子向量{x_m}。 - 对这些
D_sub维的子向量运行 K-means 算法，得到k*个质心，构成码本C_m。
- 收集所有样本的第
编码过程：
当需要编码一个新的向量 x 时：
- 将
x切分成M个子向量x_1, ..., x_M。 - 对于每个子向量
x_m：- 在对应的码本
C_m中找到距离x_m最近的质心c_{m, j_m}。 j_m = argmin_{j=1,...,k*} ||x_m - c_{m,j}||^2- 记录下这个最近质心的索引
j_m(通常是 0 到k*-1之间的整数)。
- 在对应的码本
- 将这
M个索引(j_1, j_2, ..., j_M)拼接起来，就构成了向量x的 PQ 编码。这个编码通常表示为一个字节数组（如果k*=256）。
PQcode(x) = (j_1, j_2, ..., j_M)
这个编码非常紧凑。如果 M=8，k*=256，那么无论原始向量 D 是多少维（比如 128, 256,甚至 1024），它的 PQ 编码都只需要 M * log2(k*) = 8 * 8 = 64 比特，也就是 8 个字节！相比于存储原始的 float32 向量（例如 128 维需要 128 * 4 = 512 字节），内存压缩效果非常显著。
向量 x 的量化（近似）表示 x̃ 可以由其编码对应的质心拼接而成：x̃ = [c_{1, j_1} | c_{2, j_2} | ... | c_{M, j_M}]
这个 x̃ 就是原始向量 x 在 PQ 量化后的近似重构。
PQ 的距离计算：数学推导
现在到了最关键的部分：我们只有查询向量 q 和数据库中某个向量 y 的 PQ 编码 PQcode(y) = (j_1, ..., j_M) 以及所有子空间的码本 C_1, ..., C_M，如何近似计算 q 和 y 之间的（欧氏距离平方）||q - y||^2 呢？
Faiss 主要使用了两种方法：非对称距离计算 (Asymmetric Distance Computation, ADC) 和对称距离计算 (Symmetric Distance Computation, SDC)。
1. 非对称距离计算 (ADC)
ADC 的思想是用原始查询向量 q 和量化后的数据库向量 ỹ 之间的距离来近似 ||q - y||^2。
||q - y||^2 ≈ ||q - ỹ||^2
其中 ỹ = [c_{1, j_1} | c_{2, j_2} | ... | c_{M, j_M}] 是 y 的量化重构。
我们将 q 也进行同样的子空间分解：q = [q_1 | q_2 | ... | q_M]。
那么：||q - ỹ||^2 = || [q_1 | ... | q_M] - [c_{1, j_1} | ... | c_{M, j_M}] ||^2
由于欧氏距离的平方等于各维度差值平方和，并且我们的分解是沿着维度进行的，所以这个平方距离可以分解到各个子空间上：
||q - ỹ||^2 = Σ_{m=1}^{M} ||q_m - c_{m, j_m}||^2
这个公式是 ADC 的核心！它告诉我们，计算 q 和 ỹ 的距离，等价于计算 q 的每个子向量 q_m 与 y 对应编码 j_m 所指向的质心 c_{m, j_m} 之间的距离平方，然后将这 M 个子距离加起来。
计算过程：
- 预计算/准备：
- 拥有所有码本
C_1, ..., C_M。 - 对于给定的查询向量
q，将其分解为q_1, ..., q_M。
- 拥有所有码本
- 查询时：
- 对于数据库中的每个 PQ 编码
(j_1, ..., j_M)：- 初始化总距离
d = 0。 - 对于每个子空间
m从 1 到M：- 取出查询子向量
q_m。 - 从码本
C_m中取出编码j_m对应的质心c_{m, j_m}。 - 计算子距离平方
d_m = ||q_m - c_{m, j_m}||^2。 - 累加到总距离：
d = d + d_m。
- 取出查询子向量
- 最终得到的
d就是q与该编码对应向量y的近似距离平方||q - ỹ||^2。
- 初始化总距离
- 对于数据库中的每个 PQ 编码
优化： 这个计算过程还可以进一步优化。注意到对于一个固定的查询 q，它的子向量 q_1, ..., q_M 是不变的。在遍历数据库中的编码时，我们需要反复计算 ||q_m - c_{m, j_m}||^2。我们可以预先计算出 q 的每个子向量 q_m 到其对应码本 C_m 中所有 k* 个质心的距离平方！
创建一个查找表 (Lookup Table, LUT)：LUT_m[j] = ||q_m - c_{m,j}||^2，其中 m = 1, ..., M，j = 1, ..., k*。
这个查找表的大小是 M * k* 个浮点数。
优化后的 ADC 计算过程：
- 预计算/准备：
- 拥有所有码本
C_1, ..., C_M。
- 拥有所有码本
- 查询时：
- 对于给定的查询向量
q：- 将其分解为
q_1, ..., q_M。 - 构建查找表：对于每个
m从 1 到M，计算LUT_m[j] = ||q_m - c_{m,j}||^2对于所有j = 1, ..., k*。
- 将其分解为
- 对于数据库中的每个 PQ 编码
(j_1, ..., j_M)：- 初始化总距离
d = 0。 - 对于每个子空间
m从 1 到M：- 从预计算的查找表
LUT_m中直接取出距离：d_m = LUT_m[j_m]。 - 累加到总距离：
d = d + d_m。
- 从预计算的查找表
- 最终得到的
d就是近似距离平方||q - ỹ||^2。
- 初始化总距离
- 对于给定的查询向量
通过预计算查找表，每次处理一个数据库编码时，我们只需要进行 M 次查表和 M-1 次加法！这比原来计算 M 个 D_sub 维向量的距离要快得多，尤其是当 D_sub 不算太小的时候。主要开销转移到了查询时构建查找表的那一步，但这一步对于每个查询只需要做一次。
2. 对称距离计算 (SDC)
SDC 则更进一步，它用量化后的查询向量 q̃ 和量化后的数据库向量 ỹ 之间的距离来近似 ||q - y||^2。
||q - y||^2 ≈ ||q̃ - ỹ||^2
其中 q̃ = [c_{1, i_1} | ... | c_{M, i_M}] 是 q 的量化重构（i_m 是 q_m 的最近邻质心索引），ỹ = [c_{1, j_1} | ... | c_{M, j_M}] 是 y 的量化重构。
同样地，我们可以将距离分解到子空间：
||q̃ - ỹ||^2 = Σ_{m=1}^{M} ||c_{m, i_m} - c_{m, j_m}||^2
这个公式看起来和 ADC 很像，但关键区别在于，这里的距离是两个质心之间的距离，而不是原始子向量和质心之间的距离。
SDC 的巨大优势在于： 质心之间的距离 ||c_{m, i} - c_{m, j}||^2 不依赖于具体的查询向量 q 或数据库向量 y！ 它们只依赖于码本本身。
这意味着我们可以完全预计算出所有码本内，任意两个质心之间的距离平方！
对于每个子空间 m，我们可以预计算一个 k* x k* 的距离矩阵 Dist_m，其中：Dist_m[i][j] = ||c_{m, i} - c_{m, j}||^2
这个预计算可以在索引构建完成后、任何查询发生之前就完成。总共需要存储 M * k* * k* 个距离值。
SDC 计算过程：
- 预计算/准备：
- 拥有所有码本
C_1, ..., C_M。 - 完全预计算所有子空间的质心间距离矩阵
Dist_1, ..., Dist_M。
- 拥有所有码本
- 查询时：
- 对于给定的查询向量
q：- 计算它的 PQ 编码
PQcode(q) = (i_1, ..., i_M)。
- 计算它的 PQ 编码
- 对于数据库中的每个 PQ 编码
PQcode(y) = (j_1, ..., j_M)：- 初始化总距离
d = 0。 - 对于每个子空间
m从 1 到M：- 从预计算的距离矩阵
Dist_m中直接查找：d_m = Dist_m[i_m][j_m]。 - 累加到总距离：
d = d + d_m。
- 从预计算的距离矩阵
- 最终得到的
d就是近似距离平方||q̃ - ỹ||^2。
- 初始化总距离
- 对于给定的查询向量
SDC 的查询过程极其高效！对于每个数据库编码，只需要：
- 获取查询编码
(i_1, ..., i_M)(每个查询只需计算一次)。 - 进行
M次查表（二维数组查找）和M-1次加法。
这比 ADC 的查询时计算更快，因为它连构建查询特定查找表 LUT_m 的步骤都省掉了。代价是需要额外的存储空间来存放预计算的 M * k* * k* 个质心间距离。
ADC vs SDC 对比
- 精度： 通常认为 ADC 的精度略高于 SDC。因为 ADC 使用了原始的查询向量
q，保留了更多信息，而 SDC 对查询向量q也进行了量化，引入了额外的量化误差。||q - ỹ||^2通常比||q̃ - ỹ||^2更接近真实的||q - y||^2。 - 速度： SDC 的查询速度通常快于 ADC，因为它进行了更彻底的预计算，查询时每个数据库向量的处理开销更小。
- 内存： SDC 需要额外的
M * k* * k*存储空间来存放预计算的质心间距离。ADC 则不需要这部分额外存储，但查询时需要临时空间计算M * k*大小的查找表。
在 Faiss 中，IndexPQ 对象默认使用 ADC 进行距离计算 (d = faiss.pairwise_distance(q, y_reconstructed))。但当你将 PQ 作为其他索引（如 IndexIVFPQ）的组成部分时，为了加速在倒排列表中的距离计算，Faiss 通常会利用类似 SDC 的思想（或者更准确地说，是 ADC 的预计算查找表优化）来快速筛选候选项。
进一步思考：量化误差
PQ 是一种有损压缩，量化过程必然会引入误差。误差主要来源于两个方面：
- 子空间分解： 假设子空间之间是独立的，忽略了维度间的相关性。
- K-means 量化： 每个子向量
x_m被其最近的质心c_{m, j_m}替代，||x_m - c_{m, j_m}||^2就是该子空间的量化误差。
总的量化误差（原始向量与重构向量的距离平方）是：||x - x̃||^2 = Σ_{m=1}^{M} ||x_m - c_{m, j_m}||^2
这个误差的大小直接影响了 ADC 和 SDC 计算出的距离与真实距离 ||q - y||^2 的偏差，从而影响最终的搜索精度 (Recall)。
选择合适的 M（子空间数量）和 k*（每个子空间的质心数量）是在压缩率、内存占用、计算速度和搜索精度之间进行权衡的关键。
- 增加
M(在总比特数M * log2(k*)不变的情况下，意味着k*减小，或者D_sub减小)：- 优点：编码更短（如果
k*固定），或者子空间维度D_sub更低，K-means 训练可能更有效。 - 缺点：子空间划分更细，可能破坏维度间的结构；
k*减小导致每个子空间的量化误差增大。
- 优点：编码更短（如果
- 增加
k*(通常固定为 256)：- 优点：每个子空间的量化更精细，量化误差减小。
- 缺点：编码变长（如果
M固定），码本更大，训练和存储开销增加。ADC 的查找表和 SDC 的预计算距离矩阵也更大。
实践中，k*=256 是一个非常常用的选择，因为它正好对应一个字节，便于存储和处理。M 的选择则取决于原始维度 D 和期望的压缩率/精度平衡点，通常 M 会是 8, 16, 32, 64 等。
总结
我们深入探讨了 Faiss 中 PQ 算法的数学核心：
- 原理： 将向量分解为
M个子向量，对每个子空间独立使用 K-means 学习包含k*个质心的码本。 - 编码： 将向量表示为
M个子空间最近邻质心的索引(j_1, ..., j_M)。 - 距离计算 (ADC)： 近似
||q - y||^2为||q - ỹ||^2 = Σ ||q_m - c_{m, j_m}||^2。通过预计算查询q的子向量到各码本的距离查找表LUT_m[j] = ||q_m - c_{m,j}||^2来加速。 - 距离计算 (SDC)： 近似
||q - y||^2为||q̃ - ỹ||^2 = Σ ||c_{m, i_m} - c_{m, j_m}||^2。通过完全预计算质心间距离矩阵Dist_m[i][j] = ||c_{m, i} - c_{m, j}||^2来实现极速查询。
理解这些数学推导有助于你更好地把握 PQ 的工作机制、性能特点以及精度与效率的权衡。虽然 Faiss 库为我们封装了复杂的实现细节，但洞悉其背后的数学原理，无疑能让你在应用和调优 Faiss 时更加得心应手。
希望这次的数学之旅对你有所帮助！如果你对 Faiss 或向量量化的其他方面还有疑问，欢迎继续探索和讨论。
