大规模语义检索的近似算法Product Quantization 和 IVF 的实现原理一、深度引言与场景痛点向量检索系统初期总是美好的。几千条文档、几百维向量用暴力计算做精确的 Top-K延迟稳稳控制在几十毫秒。然后业务增长文档量从几千涨到几万、几十万、几百万。暴力计算的时间线性增长延迟从 50ms 到 500ms 到 5 秒——系统就这样在你眼皮底下慢慢热死了。问题的本质是最近邻搜索的计算复杂度。给定一个查询向量 q 和 N 个数据库向量精确的 Top-K 需要计算 q 和所有 N 个向量的距离复杂度 O(N×D)D 是向量维度。当 N100 万、D1536 时一次检索要进行 15 亿次浮点运算。即便用 GPU 加速延迟也在百毫秒级别对于需要实时响应的 RAG 系统来说是不可接受的。这就是近似最近邻搜索ANN存在的意义用可控的精度损失换取数量级的性能提升。Product QuantizationPQ和 Inverted File IndexIVF是两种最经典的近似算法几乎所有现代向量数据库Milvus、Faiss、Pinecone的索引实现都建立在它们之上。二、底层机制与原理深度剖析PQ 和 IVF 解决的是不同维度的问题但组合起来威力巨大。PQ 负责压缩——把高维向量压成紧凑码大幅减少内存占用和距离计算量。IVF 负责剪枝——把搜索空间从全量缩小到相关的那一小撮。PQ 的核心思想是把每个 D 维向量切成 S 段每段 dD/S 维。对每个子空间单独做 K-Means 聚类生成 C 个聚类中心码字用码本的索引号uint8替代原始的 float32 值。原来一个向量要存 D×4 字节float32PQ 编码后只存 S×1 字节uint8 索引压缩比约 4×C/d。更重要的是距离计算——q到数据库中某向量的距离可以通过查表近似计算先算q各子段到 C 个码字的距离建一个 S×C 的查找表然后对每个数据库向量做 S 次查表和累加。复杂度从 O(N×D) 降到 O(N×S S×C×d)当 S D 时速度提升非常显著。IVF 的核心思想是在检索前先做粗聚类。构建阶段用 K-Means 把所有向量聚成 M 个簇cluster每个簇维护一个倒排列表inverted list记录属于该簇的向量。检索时先算查询向量到 M 个聚类中心的距离只去最近的 P 个nprobe簇里做精确检索。如果 P M候选集就从 N 缩到了约 P×(N/M)搜索空间大幅减小。nprobe 是速度-精度的调节旋钮nprobe 越大越接近全量检索精度高、速度慢越小越快但精度越低。PQ IVF 的组合先 IVF 粗筛缩小范围再在候选集上用 PQ 做快速距离近似。两个近似叠加精度损失可控Recall10 通常在 95%-99%速度提升可以达到 10-100 倍。这就是为什么 Faiss 的IndexIVFPQ是生产环境中最常用的索引类型。三、生产级代码实现import logging from dataclasses import dataclass from typing import Any import numpy as np logger logging.getLogger(__name__) dataclass class PQCodec: Product Quantization 编码器 n_subspaces: int # S子空间数 n_codewords: int # C每个子空间的码字数 sub_dim: int # 每个子空间的维度 codebooks: list[np.ndarray] | None None # S 个码本每个形状 (C, sub_dim) class ProductQuantizer: Product Quantization 实现 def __init__(self, dim: int, n_subspaces: int 32, n_codewords: int 256) - None: if dim % n_subspaces ! 0: raise ValueError(fdim ({dim}) must be divisible by n_subspaces ({n_subspaces})) self.dim dim self.n_subspaces n_subspaces self.n_codewords n_codewords self.sub_dim dim // n_subspaces self.codec PQCodec( n_subspacesn_subspaces, n_codewordsn_codewords, sub_dimself.sub_dim, ) def train(self, vectors: np.ndarray) - None: 训练码本每个子空间独立做 K-Means n vectors.shape[0] codebooks: list[np.ndarray] [] for s in range(self.n_subspaces): start s * self.sub_dim end start self.sub_dim sub_vectors vectors[:, start:end] # 随机选 C 个点作为初始聚类中心生产环境用 K-Means indices np.random.choice(n, min(self.n_codewords, n), replaceFalse) centroids sub_vectors[indices].copy() # 简化版 K-Means10 次迭代 for _ in range(10): # 分配点到最近的中心 dists np.sum((sub_vectors[:, np.newaxis, :] - centroids[np.newaxis, :, :]) ** 2, axis2) assignments np.argmin(dists, axis1) # 更新中心 for c in range(self.n_codewords): members sub_vectors[assignments c] if len(members) 0: centroids[c] members.mean(axis0) codebooks.append(centroids) self.codec.codebooks codebooks def encode(self, vectors: np.ndarray) - np.ndarray: 将向量编码为 PQ 码uint8 索引数组 if self.codec.codebooks is None: raise RuntimeError(Codec not trained. Call train() first.) n vectors.shape[0] codes np.zeros((n, self.n_subspaces), dtypenp.uint8) for s in range(self.n_subspaces): start s * self.sub_dim end start self.sub_dim sub_vectors vectors[:, start:end] codebook self.codec.codebooks[s] dists np.sum( (sub_vectors[:, np.newaxis, :] - codebook[np.newaxis, :, :]) ** 2, axis2, ) codes[:, s] np.argmin(dists, axis1).astype(np.uint8) return codes def decode(self, codes: np.ndarray) - np.ndarray: 从 PQ 码解码回近似向量 if self.codec.codebooks is None: raise RuntimeError(Codec not trained.) n codes.shape[0] vectors np.zeros((n, self.dim), dtypenp.float32) for s in range(self.n_subspaces): start s * self.sub_dim end start self.sub_dim vectors[:, start:end] self.codec.codebooks[s][codes[:, s]] return vectors class IVFIndex: IVF 倒排索引 def __init__(self, dim: int, n_clusters: int 1024) - None: self.dim dim self.n_clusters n_clusters self.centroids: np.ndarray | None None # 倒排列表: cluster_id - [vector_indices] self.inverted_lists: list[list[int]] [] def train(self, vectors: np.ndarray) - None: K-Means 聚类生成簇中心 n vectors.shape[0] indices np.random.choice(n, self.n_clusters, replaceFalse) self.centroids vectors[indices].copy().astype(np.float32) for _ in range(15): dists np.sum( (vectors[:, np.newaxis, :] - self.centroids[np.newaxis, :, :]) ** 2, axis2, ) assignments np.argmin(dists, axis1) for c in range(self.n_clusters): members vectors[assignments c] if len(members) 0: self.centroids[c] members.mean(axis0) def build(self, vectors: np.ndarray) - None: 建立倒排列表 self.inverted_lists [[] for _ in range(self.n_clusters)] dists np.sum( (vectors[:, np.newaxis, :] - self.centroids[np.newaxis, :, :]) ** 2, axis2, ) assignments np.argmin(dists, axis1) for i, cluster_id in enumerate(assignments): self.inverted_lists[int(cluster_id)].append(i) def search( self, query: np.ndarray, vectors: np.ndarray, k: int 10, nprobe: int 16 ) - tuple[np.ndarray, np.ndarray]: 检索 Top-K if self.centroids is None: raise RuntimeError(Index not trained.) # 1. 找最近的 nprobe 个聚类 cluster_dists np.sum((self.centroids - query) ** 2, axis1) nearest_clusters np.argpartition(cluster_dists, nprobe)[:nprobe] # 2. 收集候选向量 candidate_ids: list[int] [] for c_id in nearest_clusters: candidate_ids.extend(self.inverted_lists[int(c_id)]) if not candidate_ids: return np.array([], dtypenp.int64), np.array([], dtypenp.float32) candidate_ids np.array(candidate_ids, dtypenp.int64) candidate_vectors vectors[candidate_ids] # 3. 精确距离计算 dists np.sum((candidate_vectors - query) ** 2, axis1) # 4. Top-K top_k min(k, len(dists)) top_indices np.argpartition(dists, top_k)[:top_k] top_indices top_indices[np.argsort(dists[top_indices])] return candidate_ids[top_indices], dists[top_indices] property def avg_list_size(self) - float: 平均倒排列表长度用于评估索引质量 if not self.inverted_lists: return 0 return sum(len(lst) for lst in self.inverted_lists) / len(self.inverted_lists) def evaluate_recall( pq: ProductQuantizer, ivf: IVFIndex, vectors: np.ndarray, test_queries: np.ndarray, k: int 10, ) - dict[str, float]: 评估 PQ IVF 的召回率 exact_recalls 0 approx_recalls 0 for query in test_queries: # 精确检索ground truth exact_dists np.sum((vectors - query) ** 2, axis1) exact_top np.argpartition(exact_dists, k)[:k] # PQ IVF 近似检索 approx_ids, _ ivf.search(query, vectors, kk, nprobe32) exact_recalls len(set(exact_top) set(approx_ids)) approx_recalls len(approx_ids) return { recallk: exact_recalls / (k * len(test_queries)) if test_queries.shape[0] 0 else 0, } def main() - None: np.random.seed(42) n_vectors, dim 10000, 128 vectors np.random.randn(n_vectors, dim).astype(np.float32) # 归一化余弦相似度检索需要 vectors vectors / np.linalg.norm(vectors, axis1, keepdimsTrue) # PQ 训练和编码 pq ProductQuantizer(dimdim, n_subspaces16, n_codewords256) pq.train(vectors) codes pq.encode(vectors) decoded pq.decode(codes) # 量化误差 mse np.mean((vectors - decoded) ** 2) print(fPQ 量化误差 (MSE): {mse:.6f}) # IVF 训练和检索 ivf IVFIndex(dimdim, n_clusters100) ivf.train(vectors) ivf.build(vectors) print(fIVF 平均倒排列表大小: {ivf.avg_list_size:.1f}) # 召回率评估 test_queries np.random.randn(100, dim).astype(np.float32) test_queries test_queries / np.linalg.norm(test_queries, axis1, keepdimsTrue) recall evaluate_recall(pq, ivf, vectors, test_queries, k10) print(fRecall10: {recall[recallk]:.3f}) if __name__ __main__: logging.basicConfig(levellogging.INFO) main()代码中值得关注的点PQ 训练时每个子空间独立做 K-Means这意味着码本之间没有交互。PQ 编码把 float32 压缩为 uint8256 个码字只占 1 字节内存从 512 字节/向量128 维 × 4 字节压缩到 16 字节/向量16 个子空间 × 1 字节压缩比 32 倍。IVF 的nprobe是速度和精度的直接调节旋钮。avg_list_size可以检测聚类是否均匀——如果标准差过大说明数据分布不均需要调整 n_clusters。四、边界分析与架构权衡PQ 的精度损失主要来自量化误差——用离散码字替代连续向量值带来的信息丢失。子空间数 S 越大每段越短量化越精细但编码也越长更多字节。子空间数 S 越小每段越长编码越短但每段的信息损失越大。这里有个反直觉的现象D 维向量切分成 S 段时每段的维度 dD/S。d 越大K-Means 在 d 维空间里的聚类质量越差维度诅咒量化误差反而增大。所以 S 不是越大越好需要根据数据分布做实验确定。IVF 的精度损失来自检索漏探——如果目标向量的 true nearest neighbor 落在 nprobe 之外的簇里它就永远被遗漏了。nprobe 太小漏掉关键结果nprobe 太大退化为近似全量检索。一个实用的策略是设置nprobe sqrt(n_clusters)兼顾速度和精度。PQ IVF 组合使用时两个近似的误差会叠加。但实际测试表明两者有很好的互补性PQ 的误差偏向低频分量大方向对细节不准IVF 的误差偏向遗漏某些簇搜不到叠加后 Recall10 通常能保持在 95% 以上对于 RAG 场景完全够用。适用场景的判断如果向量库小于 10 万条暴力精确检索就够了GPU 加速下 10 万向量检索 10ms。如果 10 万到 500 万之间IVF 加精确距离计算是最佳平衡。如果超过 500 万或内存紧张上 PQ IVF 组合。如果超过 1 亿还需要考虑 HNSW 图索引或分布式分片方案。五、总结PQ 和 IVF 是大规模向量检索的两大基石算法。PQ 用子空间量化和码本索引将向量压缩 20-50 倍大幅节省内存和计算。IVF 用粗聚类将搜索空间缩小到原来的几分之一到几十分之一。组合使用实现了在百万至亿级向量规模下的毫秒级检索延迟精度损失控制在 5% 以内。选择 nprobe 和 n_subspaces 这两个参数时记住它们分别控制搜多少和存多准通过小样本 A/B 测试找到适合业务数据的 sweet spot。