乘积量化 PQ:子空间分解与短码压缩
乘积量化 PQ:子空间分解与短码压缩
在大规模向量检索任务中,高维向量带来的存储开销和距离计算成本往往令人望而生畏。乘积量化(Product Quantization,PQ)通过将高维空间分解为多个低维子空间并分别量化,以极小的精度损失换取显著的存储压缩和计算加速,是当前工业界近似最近邻(ANN)搜索的核心技术之一。本文从原理、实现到应用,带你全面掌握 PQ。
背景:向量检索的困境
在许多 AI 应用中,对象被表示为高维嵌入向量(如 512 维、768 维甚至更高)。要找到与查询最相似的向量,需要遍历全量数据计算距离,这在数据量达到百万、千万级别时面临两大瓶颈:
- 存储压力:每个向量占用的字节数大,全量存储需要巨量内存或磁盘。
- 计算压力:每次查询需与所有向量进行浮点距离运算,延迟难以接受。
近似最近邻搜索通过牺牲少量精度换取速度,而乘积量化正是其中负责压缩向量、加速距离计算的基础技术。
什么是乘积量化?
乘积量化是一种有损向量压缩方法,它把高维向量空间划分成多个互不相交的低维子空间,在每个子空间内独立地进行向量量化(聚类),然后将原始向量表示为各个子空间聚类中心的索引序列。这样一来,一个高维浮点向量就被压缩为一条短整数码,极大缩减了存储,同时可以通过查表快速计算近似距离。
通俗理解:把一个长向量切成若干段,每段都用最接近的“代表向量”来替代,最后用这些代表向量的编号拼起来作为压缩编码。
核心原理拆解
1. 子空间分解
给定一个 D 维向量 x ∈ R^D,我们将其分割为 m 个子向量,每个子向量的维度为 d = D / m(通常要求整除)。例如:
- 原始向量维度
D=128,设定m=8,则每个子空间维度d=16。 - 子向量
x = [x_1 | x_2 | … | x_m],其中x_j ∈ R^d。
这种分割的本质是将原始的高维空间笛卡尔积分解为 m 个低维空间的乘积,故此得名“乘积”量化。
2. 子空间独立量化
在每个子空间内,我们独立训练一个码本。码本 C_j 包含 k 个代表向量(即聚类中心),通常通过 k-means 算法得到。C_j = {c_{j,1}, c_{j,2}, …, c_{j,k}},每个 c 是一个 d 维向量。
对第 j 段子向量 x_j,量化过程是将其映射到距离最近的码本向量:
q_j(x) = arg min_{c ∈ C_j} ||x_j - c||
该映射的结果就是码本中的索引 i ∈ {1,…,k}。
3. 短码压缩
完成所有子空间的量化后,原始向量 x 被表示为 m 个索引的元组:
code(x) = (i_1, i_2, …, i_m),每个 i_j ∈ [0, k-1]
每个索引只需 log2(k) 比特存储。总存储量为 m * log2(k) 比特,远小于原始 D * 32 比特(单精度浮点)。例如 D=128,m=8,k=256,压缩后每条向量仅需 8 * 8 = 64 bits = 8 bytes,相比原始的 128*4=512 bytes 压缩了 64 倍。
4. 距离计算:非对称距离计算(ADC)
PQ 最巧妙之处在于无需重建完整向量即可近似计算查询向量与数据库向量的距离,主要通过**非对称距离计算(Asymmetric Distance Computation,ADC)**实现。
对于查询向量 q,同样将其分割为 q_1, …, q_m。为计算 q 与数据库中某向量 x(其编码已知)的近似距离:
- 预先计算查询子向量到每个子空间码本所有中心向量的距离表:
d_{j,i} = ||q_j - c_{j,i}||^2,共m * k个距离值。 - 对于编码为
(i_1, …, i_m)的数据库向量,其近似距离通过查表累加得到:
dist(q, x) ≈ Σ_{j=1}^{m} d_{j, i_j}
该过程将原本 O(D) 的浮点距离计算转化为 m 次查表加法,速度提升显著。注意这里查询端不量化(非对称),而是利用数据库端的码本索引,故称“非对称”。
另一种是对称距离计算(SDC),即查询端也量化后再计算距离,但由于额外引入查询量化误差,精度通常不如 ADC,因此实践中几乎都使用 ADC。
训练与编码流程
构建一个 PQ 索引通常包含以下步骤:
- 确定参数:选择子空间数量
m和每个子空间的码本大小k。D需被m整除。 - 采样训练数据:从待索引向量中随机采样足够数量(如 10 万)作为训练集。
- 划分子空间:将所有训练向量的各段子向量拆分出来。
- 逐空间聚类:对第
j个子空间的全部子向量运行 k-means(k个中心),存储中心向量作为码本C_j。 - 编码数据库:将数据库中的每个向量分块,于对应子空间内找到最近中心,记录索引编号。最终数据库只存储整数编码和对应码本。
参数选择建议
m(子空间数):m越大,压缩越细粒度,通常编码长度越长,精度保留越好,但计算距离时的查表加法次数增多。需权衡存储和精度。常见值:m = D/2到D的一半左右,但要保证d不至于太小(一般d ≥ 1即可,但过小量化误差大)。k(子空间中心数):k通常取 256(8 bits),这样每个索引占用 1 字节,计算方便。若使用 2^16=65536 则每索引 16 bits,精度更高但存储和距离表更大。典型配置k=256是速度与精度的平衡点。d = D / m:为子空间维度。注意 k-means 在维度d较高时量化误差会增大,所以d一般不超过 16~32 为宜。如果原始维度很高,需相应增大m以减小每个子空间的维度。
优势与局限
优点
- 极致压缩:将高维浮点向量压缩为每向量几十到几百字节,非常适合 billion 级数据的内存存储。
- 快速距离计算:ADC 将耗时的浮点运算转换为整型查表加法,通常可以有数量级的加速。
- 实现简单:核心仅依赖于 k-means,工程落地成熟。
局限
- 量化误差:有损压缩会损失距离精度,尤其当子空间维度过高或码本太小时会比较明显。
- 数据集依赖:码本质量严重依赖训练数据的分布,对于与训练集差异较大的新数据,误差可能增大。
- 固定长度编码:标准 PQ 为所有向量分配相同的编码长度,无法自适应调整。
工程实践中的增强变体
IVFPQ:倒排索引 + PQ
单独使用 PQ 仍需要遍历所有编码进行距离累加,数据量超大时线性扫描依然较慢。最常见的是将**倒排索引(IVF)**与 PQ 结合,即使用 IVF 对向量空间进行粗分桶,每个桶内的向量再使用 PQ 编码。检索时,先通过粗量化找到最近的若干桶,仅对这些桶内的 PQ 编码进行 ADC 重排。这就是 Faiss 中的 IndexIVFPQ,是生产环境中最常用的方案。
OPQ:优化乘积量化
原始 PQ 只是简单地将向量维度顺序切分,未考虑维度间的相关性。优化乘积量化(OPQ)通过一个正交旋转矩阵对向量进行预处理,使得能量在各个子空间均匀分布或类内相关性降低,大幅提升量化精度。Faiss 中对应 IndexOPQ。
加性和积量化(AQ / LSQ)
采用多级残差编码替代单一的 M 分子空间分解,在同等码长下通常能获得更高的召回率,但训练和编码更复杂。
应用场景
- 图像检索:对深度卷积网络提取的图片 embedding 进行 PQ 压缩,实现以图搜图与大规模图片去重。
- 推荐系统:压缩用户和物品的 embedding,加速召回阶段的相似度计算。
- 自然语言处理:对句子或文档的稠密向量进行压缩,支撑语义搜索。
- 生物信息学:对蛋白质或基因序列的高维特征做快速比对。
总结
乘积量化 PQ 是一种优雅而实用的向量压缩技术,通过子空间分解将高维空间拆分为低维子空间的笛卡尔积,再在各子空间内独立量化,从而将向量表示为短整数短码。配合非对称距离计算(ADC),可以在不显著损失检索精度的情况下,实现存储和计算的双重压缩。掌握 PQ 的原理和参数配置,是构建大规模向量检索引擎的重要一步。
对于入门实践,推荐从 Faiss 库的 IndexPQ 和 IndexIVFPQ 入手,直观感受不同参数对召回和速度的影响,进而在自身业务数据上寻找最优配置。