资讯动态

向量乘积量化 PQ 内核剖析:高维子空间切分与码本查找表聚类原理

发布时间:2026/10/15 12:30:12 来源:尧图企业网站定制
在大规模向量检索系统中海量高维浮点向量的存储成本与检索吞吐构成了最尖锐的物理矛盾。以常见开源嵌入模型输出的 1024 维单精度浮点FP32向量为例单条向量占用物理内存高达 4096 字节。当系统承载十亿级别向量检索底库时纯向量数据便需吞噬超过 4TB 的裸内存空间。若叠加 HNSW 等图索引的邻居指针开销物理节点集群的硬件物料成本BOM与维护总拥有成本TCO将彻底击穿业务预算。从存储架构师的财务账本来看纯粹依靠堆叠 DRAM 内存来承载稠密向量检索并不具备工程可持续性。乘积量化Product Quantization, 简称 PQ是破除这堵“内存墙”的核心手段。通过牺牲微小的召回精度乘积量化可将 1024 维 FP32 向量压缩至 64 字节乃至 32 字节压缩比高达 64 倍到 128 倍将原本需要数十台高配内存型服务器承载的底库压缩收敛至单台配备常规 NVMe 缓存的机器中。高维子空间正交切分与质心码本构造乘积量化的数学本质是将极高维度的欧氏几何空间解耦切分为若干个相互正交的低维子空间并对每个子空间实施独立的向量量化Vector Quantization。设原始向量空间为 $D$ 维欧氏空间 $\mathbb{R}^D$。算法将其均匀切分为 $M$ 个连续且互不重叠的子空间每个子空间的维度为 $d D / M$。对于底库中的任意向量 $x \in \mathbb{R}^D$可将其正交投影表示为 $M$ 个子向量的笛卡尔积连接$$x [u_1(x), u_2(x), \dots, u_M(x)], \quad u_m(x) \in \mathbb{R}^{d}$$完成空间切分后系统在每个子空间内独立运行 K-Means 聚类算法。在工业落地中每个子空间通常选定聚类中心数 $K 256$。这一取值具备极其考究的系统工程意义256 个聚类中心刚好能用无符号 8 位整型1 个 Byte即 uint8完整索引。这 $K$ 个聚类中心被称为该子空间的“码字”Codeword$K$ 个码字构成了子空间码本Codebook$C_m {c_{m,1}, c_{m,2}, \dots, c_{m,K}}$。当全量底库数据在各个子空间中完成聚类后对于原始向量 $x$ 在第 $m$ 个子空间上的切片 $u_m(x)$只需计算其与该子空间内 256 个质心的欧氏距离找出欧氏距离最近的质心编号 $k \in [0, 255]$$$q_m(x) \arg\min_{k \in [0, 255]} |u_m(x) - c_{m,k}|^2$$此时原始占用 $D \times 4$ 字节的浮点向量被彻底压缩为一个长度为 $M$ 字节的 uint8 数组 $[q_1(x), q_2(x), \dots, q_M(x)]$。原本 1024 维浮点向量4096 字节若切分为 $M 64$ 个 16 维子空间压缩后仅占用 64 字节物理压缩比达到 98.44%。距离计算范式对称距离与非对称距离的性能与精度边界压缩后的向量如何进行高效相似度检索乘积量化派生出了两种度量模式对称距离计算Symmetric Distance Computation, SDC与非对称距离计算Asymmetric Distance Computation, ADC。对称距离计算要求查询向量Query同样经过子空间切分与码本量化转为 $M$ 字节的码字索引然后通过查阅预先计算好的码本间距离矩阵来计算得分。这种模式引入了“双重空间量化误差”检索召回率下滑严重在严谨的生产系统中基本已被弃用。现代向量检索引擎全面采用非对称距离计算ADC。ADC 的核心设计原则是查询向量 Query 坚决不进行量化保持原始高精度的 FP32 浮点状态。当查询向量 $y$ 到达引擎内核时ADC 的执行流水线分为两个确定性阶段构建预计算距离查找表Distance Lookup Table, LUT将 Query 向量 $y$ 同样切分为 $M$ 个 $d$ 维子向量 $[u_1(y), \dots, u_M(y)]$。在每个子空间 $m$ 中独立计算该子向量与码本中 256 个质心 $c_{m,k}$ 的欧氏距离平方将浮点结果填充入一张 $M \times 256$ 大小的二维查找表$$LUT[m][k] |u_m(y) - c_{m,k}|^2, \quad m \in [0, M-1], k \in [0, 255]$$构建该表的时间复杂度仅为 $O(M \times K \times d) O(D \times K)$。对于 1024 维向量、$K256$仅需执行约 26 万次浮点运算在现代 CPU 上耗时不超过几十微秒摊薄在后续海量底库比对中几乎可以忽略。底库扫描与查表聚合底库中的每一条被压缩向量本质上是一个长度为 $M$ 的 uint8 偏移量数组。计算 Query 与底库向量 $x$ 的近似距离时CPU 完全摆脱了昂贵的浮点乘加指令与开方运算退化为极速的内存查表累加操作$$\tilde{d}(y, x) \sum_{m0}^{M-1} LUT[m][q_m(x)]$$在硬件层面上长度为 $M \times 256 \times 4$ 字节的查找表以 $M64$ 为例仅占 64KB能够被完整锁死在 CPU 的 L1/L2 数据高速缓存L1/L2 Data Cache内。遍历底库向量时通过连续内存预取与 SIMD 指令并行查表累加吞吐量相较于原生 FP32 暴力检索提升了 1 到 2 个数量级。工业级纯 Python/NumPy 内核实现下面提供一个符合工业级逻辑的乘积量化完整实现包含高维空间正交切分、多子空间 K-Means 码本离线训练、向量量化压缩、以及查询端构建 LUT 查找表进行极速 ADC 距离扫描的过程。import numpy as np from sklearn.cluster import KMeans class ProductQuantizer: def __init__(self, total_dim: int, sub_quantizers: int, k_centroids: int 256): assert total_dim % sub_quantizers 0, 总维度必须能被子空间数量整除 assert k_centroids 256, 质心数不可超过 256以便使用 uint8 存储 self.total_dim total_dim self.M sub_quantizers self.d_prime total_dim // sub_quantizers self.K k_centroids # 码本矩阵 shape: (M, K, d_prime)float32 类型 self.codebooks np.zeros((self.M, self.K, self.d_prime), dtypenp.float32) self.is_trained False def fit(self, dataset: np.ndarray, max_iter: int 20): 在各个子空间独立训练 K-Means 码本 dataset shape: (N, total_dim) num_samples, dim dataset.shape assert dim self.total_dim, 数据维度与配置不一致 for m in range(self.M): start_col m * self.d_prime end_col (m 1) * self.d_prime # 切分出第 m 个子空间的训练样本集 sub_data dataset[:, start_col:end_col].astype(np.float32) # 使用 K-Means 聚类生成 256 个质心 kmeans KMeans( n_clustersself.K, max_itermax_iter, n_init3, random_state42 m ) kmeans.fit(sub_data) self.codebooks[m] kmeans.cluster_centers_.astype(np.float32) self.is_trained True def encode(self, vectors: np.ndarray) - np.ndarray: 将原始浮点向量量化压缩为 uint8 编码数组 vectors shape: (N, total_dim) 返回 shape: (N, M), dtype: np.uint8 assert self.is_trained, 码本尚未训练 num_vectors vectors.shape[0] quantized_codes np.zeros((num_vectors, self.M), dtypenp.uint8) for m in range(self.M): start_col m * self.d_prime end_col (m 1) * self.d_prime sub_vecs vectors[:, start_col:end_col].astype(np.float32) # (N, d) centroids self.codebooks[m] # (K, d) # 利用广播机制计算子向量与各个质心的欧氏距离平方 # ||x - c||^2 ||x||^2 ||c||^2 - 2 * x * c^T x_sq np.sum(sub_vecs ** 2, axis1, keepdimsTrue) # (N, 1) c_sq np.sum(centroids ** 2, axis1, keepdimsTrue).T # (1, K) dot np.dot(sub_vecs, centroids.T) # (N, K) dists x_sq c_sq - 2 * dot # 取最近质心索引 quantized_codes[:, m] np.argmin(dists, axis1).astype(np.uint8) return quantized_codes def build_lut(self, query: np.ndarray) - np.ndarray: 为未量化的单条原始 Query 向量构建距离查找表 (LUT) query shape: (total_dim,) 返回 shape: (M, K), dtype: np.float32 assert self.is_trained, 码本尚未训练 lut np.zeros((self.M, self.K), dtypenp.float32) for m in range(self.M): start_col m * self.d_prime end_col (m 1) * self.d_prime q_sub query[start_col:end_col].astype(np.float32) # (d,) centroids self.codebooks[m] # (K, d) # 计算欧氏距离平方 diff centroids - q_sub # 广播: (K, d) - (d,) - (K, d) lut[m] np.sum(diff ** 2, axis1) return lut def search_adc(self, query: np.ndarray, codes: np.ndarray, top_k: int 10): 非对称距离扫描查表累加获得候选 Top-K codes shape: (N, M), dtype: np.uint8 # 第一阶段构建查找表 (M, K) lut self.build_lut(query) # 第二阶段查表并累加每个向量在各个子空间的距离 num_candidates codes.shape[0] approx_distances np.zeros(num_candidates, dtypenp.float32) # 遍历 M 个子空间查表叠加 for m in range(self.M): sub_codes codes[:, m] # (N,) 取所有向量在第 m 子空间的质心索引 approx_distances lut[m, sub_codes] # 获取最小距离的前 Top-K 索引 if top_k num_candidates: top_k_indices np.argsort(approx_distances) else: top_k_indices np.argpartition(approx_distances, top_k)[:top_k] top_k_indices top_k_indices[np.argsort(approx_distances[top_k_indices])] return top_k_indices, approx_distances[top_k_indices] if __name__ __main__: np.random.seed(42) # 模拟 10 万条 128 维向量划分为 16 个子空间每个子空间 8 维 dim 128 sub_m 16 n_base 100000 n_train 20000 raw_base np.random.randn(n_base, dim).astype(np.float32) pq ProductQuantizer(total_dimdim, sub_quantizerssub_m, k_centroids256) print([离线训练] 开始切分子空间并聚类码本...) pq.fit(raw_base[:n_train]) print([向量编码] 开始将底库浮点向量量化压缩为 uint8 编码矩阵...) compressed_codes pq.encode(raw_base) print(f原始底库内存占用: {raw_base.nbytes / (1024 * 1024):.2f} MB) print(f压缩底库内存占用: {compressed_codes.nbytes / (1024 * 1024):.2f} MB) print(f压缩比例: {raw_base.nbytes / compressed_codes.nbytes:.1f} 倍) # 模拟单次在线查询 q np.random.randn(dim).astype(np.float32) top_indices, distances pq.search_adc(q, compressed_codes, top_k5) print(f[在线检索] ADC 查表扫描完成前 5 候选索引: {top_indices}近似距离: {distances})存储架构落地避坑与系统权衡在严苛的生产环境落地乘积量化时必须在工程设计层面规避以下物理陷阱内存连续布局与 Cache Line 伪共享防御量化后的底库矩阵在内存中必须以行主序Row-Major连续存储。对于 $M64$ 的场景每个向量恰好占用 64 字节完美契合 x86 与 ARM 架构下单个 Cache Line64B的硬件步长。若设计为列存储或指针链式结构底库扫描时将会引发极为严重的 Cache Miss使查表带来的吞吐优势彻底被内存总线争用吞噬。倒排乘积量化IVF-PQ的剪枝界限纯 PQ 虽然压缩了空间并简化了计算但其算法本质依然是 $O(N)$ 复杂度的暴力穷举扫描。在十亿级规模下必须与倒排索引Inverted File, IVF组合使用。通过一级粗量化聚类出数千至数万个倒排桶Voronoi Cells查询时仅探查最近的几十个倒排桶并在桶内执行 PQ 残差量化编码的比对。此时残差向量的方差大幅收窄既控制了搜索空间又显著补偿了量化带来的精度损失。召回率回捞的精排Rerank开销治理量化不可避免地会引入距离畸变。对于高精确度要求的金融或电商风控业务通常采用“PQ 快速召回 Top-1000 候选集 $\rightarrow$ 读取原始 FP32 向量进行精确欧氏距离重排 Top-50”的级联架构。此时原始向量不能常驻昂贵的 DRAM而应顺序追加刷入 NVMe SSD借助 Linux 系统的pread系统调用或 Direct I/O 进行点查。这一机制把昂贵的浮点计算缩减至千次以内实现了极致内存节省与严谨业务召回的工业级平衡。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑