土法炼钢兴趣小组的算法知识备份

乘积量化与 IVF-PQ:十亿向量的工业级检索

目录

你有十亿条 128 维的 float32 向量。HNSW 索引需要多少内存?

每个向量 512 字节,加上图结构每个节点大约 640 字节的开销——总计约 600 GB。即便用最贵的云实例,这也是一笔不小的账单。更现实的情况是,你的预算只有 32 GB 内存,而业务要求毫秒级响应、90% 以上的

这时候你需要的不是更大的机器,而是更聪明的压缩。

Product Quantization(乘积量化,简称 PQ)是 Herve Jegou 等人在 2011 年发表的向量压缩方法。它把 128 维向量压缩到 8 个字节,压缩比 64 倍,却仍然能在压缩域上直接计算近似距离。把 PQ 与 Inverted File(倒排索引,简称 IVF)组合,就得到了 IVF-PQ——迄今为止处理十亿级向量检索最成熟的工业方案。

Facebook 用它检索数十亿张图片的相似内容。Spotify 用它在数亿首歌的嵌入空间里做推荐。Milvus、Qdrant、Weaviate 等向量数据库,底层的大规模索引引擎几乎都绕不开 IVF-PQ。

这篇文章从向量压缩的基本问题出发,一路推进到完整的 C 实现和十亿级基准测试。

一、向量检索的内存墙

暴力搜索的代价

给定数据库 \(X = \{x_1, x_2, \ldots, x_N\}\),每个向量 \(x_i \in \mathbb{R}^D\),查询向量 \(q\),最近邻搜索需要计算 \(N\) 次距离:

\[ \text{NN}(q) = \arg\min_{x_i \in X} \|q - x_i\|^2 \]

\(N = 10^9\)\(D = 128\),每个向量 512 字节,光是存储原始向量就需要:

\[ 10^9 \times 512 \text{ bytes} = 512 \text{ GB} \]

暴力搜索在单核上遍历这些向量需要数十分钟。即便用 SIMD 加速也无法改变根本问题:数据太大了。

索引方法的内存开销

方法 存储开销(10 亿 128 维 float32) 附加索引开销
Brute-force 512 GB 0
HNSW(M=32) 512 GB ~90 GB(图结构)
IVF-Flat(nlist=65536) 512 GB ~34 MB(质心)
IVF-PQ(M=8, nlist=65536) 8 GB ~34 MB(质心)+ 2 MB(码本)

IVF-PQ 把存储从 512 GB 压到 8 GB——整整两个数量级。这就是我们要花一整篇文章理解的压缩魔法。

二、标量量化:最朴素的压缩

在进入 PQ 之前,先看最简单的压缩方案:标量量化(Scalar Quantization, SQ)。

SQ8:逐维度线性映射

思路很直接:对每一维,统计所有向量在该维度的最小值 \(v_{\min}^{(d)}\) 和最大值 \(v_{\max}^{(d)}\),然后将 float32 线性映射到 uint8:

\[ \hat{x}^{(d)} = \text{round}\left(\frac{x^{(d)} - v_{\min}^{(d)}}{v_{\max}^{(d)} - v_{\min}^{(d)}} \times 255\right) \]

还原时做逆映射:

\[ \tilde{x}^{(d)} = \hat{x}^{(d)} \times \frac{v_{\max}^{(d)} - v_{\min}^{(d)}}{255} + v_{\min}^{(d)} \]

压缩比 4 倍(float32 -> uint8),128 维向量从 512 字节压到 128 字节。

SQ4:半字节量化

更激进一点,用 4-bit 表示每个维度,可选值 0-15。两个维度共享一个字节,压缩比 8 倍。代价是量化误差显著增大——对于值域跨度大的维度,16 个桶远远不够。

SQ 的局限

标量量化独立对待每一维,完全忽略了维度之间的相关性。如果向量的能量集中在某个子空间,SQ 无法利用这一结构。更根本的问题是,SQ 的压缩比上限有限——即便用 SQ4 也只有 8 倍,对于十亿级场景还是不够。

我们需要跨维度的联合量化。

三、乘积量化:子空间分解与码本

Product Quantization 的核心思想可以用一句话概括:把高维空间拆成若干低维子空间,在每个子空间里独立做 k-means 聚类,然后用聚类编号代替原始子向量。

子空间分解

\(D\) 维向量 \(x\) 均匀切分为 \(M\) 个子向量:

\[ x = [u_1, u_2, \ldots, u_M], \quad u_m \in \mathbb{R}^{D/M} \]

例如 \(D = 128\)\(M = 8\),每个子向量 \(u_m\) 是 16 维。

码本学习

对每个子空间 \(m\),从训练集的所有子向量中运行 k-means,得到 \(K\) 个质心(centroid):

\[ C_m = \{c_{m,0}, c_{m,1}, \ldots, c_{m,K-1}\}, \quad c_{m,k} \in \mathbb{R}^{D/M} \]

通常取 \(K = 256\),这样每个子向量的编号恰好用一个字节(uint8)表示。

编码

给定向量 \(x\),PQ 编码过程是:对每个子空间 \(m\),找到最近的质心编号:

\[ q_m(x) = \arg\min_{k \in \{0, \ldots, K-1\}} \|u_m - c_{m,k}\|^2 \]

最终得到 PQ 码:

\[ \text{PQ}(x) = [q_1(x), q_2(x), \ldots, q_M(x)] \]

这是一个 \(M\) 字节的紧凑表示。对于 \(D = 128\)\(M = 8\) 的配置:

PQ 子空间分解与码本查找

为什么叫”乘积量化”

因为整体的码本是各子空间码本的笛卡尔积:

\[ C = C_1 \times C_2 \times \cdots \times C_M \]

有效码字数量为 \(K^M = 256^8 \approx 1.8 \times 10^{19}\)——一个天文数字的码本大小,但实际存储只需要 \(M \times K \times (D/M) = 8 \times 256 \times 16 = 32768\) 个 float,约 128 KB。这就是乘积结构的威力:用指数级的表达能力换取线性的存储开销。

四、非对称距离计算(ADC)

有了 PQ 码,怎么计算查询向量 \(q\) 与数据库向量 \(x\) 之间的距离?

对称距离计算(SDC)

最朴素的做法是把 \(q\) 也编码成 PQ 码,然后比较两个 PQ 码之间的距离。这叫对称距离计算(Symmetric Distance Computation):

\[ \hat{d}(q, x) = \sum_{m=1}^{M} \|c_{m,q_m(q)} - c_{m,q_m(x)}\|^2 \]

可以预计算所有 \(K \times K\) 的质心间距离,做成查找表。但问题是:查询向量本身也被量化了,引入了双重误差。

非对称距离计算(ADC)

更好的方案是保持查询向量 \(q\) 不量化,只量化数据库向量 \(x\)

\[ \tilde{d}(q, x) = \sum_{m=1}^{M} \|q_m - c_{m,q_m(x)}\|^2 \]

其中 \(q_m\) 是查询向量在第 \(m\) 个子空间的原始子向量。

关键优化:预计算查找表。在搜索开始前,对每个子空间 \(m\),计算查询子向量 \(q_m\) 到所有 \(K\) 个质心的距离:

\[ \text{LUT}[m][k] = \|q_m - c_{m,k}\|^2, \quad m \in [1, M], \; k \in [0, K) \]

查找表大小为 \(M \times K = 8 \times 256 = 2048\) 个 float,约 8 KB——完全放得进 L1 缓存。

然后对每个数据库向量的 PQ 码 \([q_1, q_2, \ldots, q_M]\),距离计算退化为 \(M\) 次查表加法:

\[ \tilde{d}(q, x) = \sum_{m=1}^{M} \text{LUT}[m][q_m(x)] \]

这就是 ADC 的核心:把昂贵的浮点向量距离计算,变成了几次整数索引查表和加法。 单个距离计算从 \(O(D)\) 的浮点乘加变成了 \(M\) 次内存访问加 \(M\) 次加法。

ADC 的误差分析

ADC 的距离估计误差来自于量化误差。对于子空间 \(m\),设量化误差 \(\epsilon_m = u_m - c_{m,q_m(x)}\),则:

\[ \tilde{d}(q, x) = \|q - x\|^2 + 2\sum_{m=1}^{M} \langle q_m - u_m, \epsilon_m \rangle + \sum_{m=1}^{M} \|\epsilon_m\|^2 \]

中间的交叉项期望为零(因为 k-means 质心是条件均值),最后一项是量化失真。因此 ADC 是真实距离的有偏估计,偏差等于总量化失真。增大 \(K\)\(M\) 都可以减小失真,但相应地增大码本或编码长度。

五、倒排文件索引(IVF)

PQ 解决了存储问题,但搜索时仍然要遍历所有 \(N\) 个 PQ 码。对于 \(N = 10^9\),即便每个距离只需要 \(M\) 次查表,总计算量仍然巨大。

IVF(Inverted File Index)通过一个粗量化器(coarse quantizer)把向量空间划分成若干 Voronoi cell,搜索时只访问少数几个 cell,从而大幅减少搜索范围。

粗量化器

训练一个 k-means 模型,将向量空间划分为 \(L\)(通常记作 nlist)个簇:

\[ \mu_1, \mu_2, \ldots, \mu_L \in \mathbb{R}^D \]

每个数据库向量被分配到最近的簇:

\[ \text{cell}(x) = \arg\min_{l \in [1, L]} \|x - \mu_l\|^2 \]

倒排列表

为每个簇维护一个倒排列表(inverted list),存储属于该簇的所有向量(或向量的 PQ 码)。

cell 0:  [x_3, x_17, x_892, ...]
cell 1:  [x_1, x_45, x_1003, ...]
cell 2:  [x_7, x_28, x_556, ...]
...
cell L-1: [x_99, x_478, ...]

搜索流程

给定查询向量 \(q\)

  1. 计算 \(q\) 到所有 \(L\) 个质心的距离,找到最近的 \(\text{nprobe}\) 个簇。
  2. 只扫描这 \(\text{nprobe}\) 个簇的倒排列表。
  3. 在扫描的向量中找到最近邻。

\(L = 65536\)\(\text{nprobe} = 64\) 时,平均只扫描 \(64 / 65536 \approx 0.1\%\) 的数据。

nprobe 的权衡

nprobe 是 IVF 最重要的运行时参数:

nprobe 扫描比例 (典型) 延迟(典型)
1 1/65536 0.35 0.1 ms
8 8/65536 0.70 0.5 ms
64 64/65536 0.90 3 ms
256 256/65536 0.97 12 ms
65536 100% 1.00 ~暴力搜索

实践中,nprobe 通常设置为 \(\sqrt{L}\)\(4\sqrt{L}\),在召回率和速度之间取一个合理的平衡。

六、IVF-PQ:十亿级检索的组合拳

把 IVF 和 PQ 组合起来,就得到了 IVF-PQ——目前处理十亿级向量最成熟的单机方案。

索引结构

IVF-PQ 的索引包含三部分:

  1. 粗量化器\(L\)\(D\) 维质心,存储为 float32,占 \(L \times D \times 4\) 字节。
  2. PQ 码本\(M\) 个子空间,每个子空间 \(K\) 个质心,占 \(M \times K \times (D/M) \times 4\) 字节。
  3. 倒排列表:每个数据库向量用 \(M\) 字节的 PQ 码存储,加上 8 字节的向量 ID,共 \((M + 8) \times N\) 字节。

对于 \(N = 10^9\)\(D = 128\)\(M = 8\)\(L = 65536\)

总计约 16 GB——一台普通的 32 GB 服务器就能装下十亿向量。

搜索流程(伪代码)

def ivfpq_search(query, index, nprobe, topk):
    # 阶段一:粗量化
    # 计算 query 到所有 L 个 coarse centroids 的距离
    coarse_dists = compute_distances(query, index.coarse_centroids)  # L 次距离
    probe_cells = argsort(coarse_dists)[:nprobe]

    # 阶段二:构建 PQ 查找表
    # 对 query 的每个子向量,计算到 K 个 PQ 质心的距离
    lut = np.zeros((M, K), dtype=np.float32)
    for m in range(M):
        q_sub = query[m * dsub : (m + 1) * dsub]
        for k in range(K):
            lut[m][k] = np.sum((q_sub - codebook[m][k]) ** 2)

    # 阶段三:扫描倒排列表
    candidates = MaxHeap(capacity=topk)
    for cell_id in probe_cells:
        for (vec_id, pq_code) in index.invlists[cell_id]:
            dist = sum(lut[m][pq_code[m]] for m in range(M))
            candidates.push(vec_id, dist)

    return candidates.to_sorted_list()

残差编码

一个重要的优化:在将向量存入倒排列表时,不直接对原始向量做 PQ 编码,而是对残差(向量减去其所属簇的质心)做 PQ 编码:

\[ r = x - \mu_{\text{cell}(x)} \]

残差的分布更紧凑(均值接近零,方差更小),PQ 能更准确地量化。这通常可以提升 2-5 个百分点的 recall。

搜索时,距离计算变为:

\[ \|q - x\|^2 = \|q - \mu_l - r\|^2 \]

其中 \(q - \mu_l\) 是查询向量相对于当前探测簇质心的残差。因此对每个探测簇,需要重新计算一次查找表——代价不大,因为 nprobe 通常只有几十。

七、OPQ:学习最优旋转矩阵

标准 PQ 有一个假设:各子空间的量化误差大致均等。但实际数据往往不满足这一条件——某些维度的方差远大于其他维度,导致某些子空间的量化误差远大于其他子空间。

Optimized Product Quantization(OPQ)通过学习一个正交旋转矩阵 \(R\) 来解决这个问题。

目标函数

OPQ 优化的目标是最小化总量化失真:

\[ \min_{R, C_1, \ldots, C_M} \sum_{i=1}^{N} \|Rx_i - \tilde{q}(Rx_i)\|^2 \]

其中 \(R \in \mathbb{R}^{D \times D}\) 是正交矩阵(\(R^TR = I\)),\(\tilde{q}\) 是 PQ 量化函数。

交替优化

OPQ 通过交替优化求解:

  1. 固定 \(R\),优化码本:这就是标准的 PQ 训练(对旋转后的向量做 k-means)。
  2. 固定码本,优化 \(R\):这是一个 Procrustes 问题,有解析解。设 \(X\) 为原始向量矩阵,\(\hat{X}\) 为量化重建向量矩阵,则最优 \(R\) 满足:

\[ R = VU^T, \quad \text{where } U\Sigma V^T = \text{SVD}(\hat{X}^T X) \]

效果

OPQ 通常能在不增加编码长度的情况下,提升 3-8 个百分点的 recall。在 Faiss 中对应 OPQx_y 前缀(如 OPQ8_64),其中 \(x\) 是子空间数,\(y\) 是旋转后的维度。

import faiss

# 标准 PQ
index_pq = faiss.index_factory(128, "PQ8")

# OPQ + PQ
index_opq = faiss.index_factory(128, "OPQ8_64,PQ8")

# OPQ + IVF-PQ
index_opq_ivfpq = faiss.index_factory(128, "OPQ8_64,IVF65536,PQ8")

八、Faiss 实现详解

Faiss 是 Meta(原 Facebook)开源的向量检索库,是 IVF-PQ 最成熟的工业级实现。

index_factory 语法

Faiss 用一个字符串描述索引结构,叫做 index_factory

import faiss
import numpy as np

D = 128          # 向量维度
N = 10_000_000   # 数据库大小
N_train = 500_000  # 训练集大小

# 生成数据(实际场景从文件加载)
xb = np.random.randn(N, D).astype('float32')
xt = xb[:N_train]   # 训练集
xq = np.random.randn(100, D).astype('float32')  # 查询集

# 构建索引
index = faiss.index_factory(D, "IVF65536,PQ8")

# 训练
index.train(xt)

# 添加向量
index.add(xb)

# 搜索
index.nprobe = 64
D_result, I_result = index.search(xq, 10)  # top-10

常用 index_factory 配置

"Flat"                      暴力搜索
"PQ8"                       纯 PQ,8 字节编码
"IVF4096,Flat"              IVF + 精确距离
"IVF65536,PQ8"              IVF-PQ,最常用配置
"IVF65536,PQ16"             更高精度的 IVF-PQ
"OPQ8_64,IVF65536,PQ8"     OPQ 旋转 + IVF-PQ
"IVF65536,PQ8+16"           PQ 粗编码 + 16 字节细化
"PCAR64,IVF65536,PQ8"      PCA 降维 + IVF-PQ

GPU 加速

Faiss 的 GPU 版本可以在单张 GPU 上训练和搜索:

# 将 CPU 索引转为 GPU 索引
res = faiss.StandardGpuResources()
gpu_index = faiss.index_cpu_to_gpu(res, 0, index)

# 多 GPU
gpu_index = faiss.index_cpu_to_all_gpus(index)

GPU 上的 k-means 训练速度是 CPU 的 10-50 倍,搜索吞吐量可以达到每秒百万级查询。

训练数据量的选择

PQ 码本训练需要多少数据?经验法则:

实际中取 \(\min(N, 10^6)\) 作为训练集通常足够。

九、C 实现:PQ 编码器与 ADC 搜索

下面给出一个完整的 C 实现,包含 PQ 码本训练(简化版)、编码和 ADC 搜索。

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <float.h>
#include <time.h>

/* ========== 配置 ========== */
#define DIM       128          /* 向量维度 */
#define M_SUB     8            /* 子空间数 */
#define DSUB      (DIM / M_SUB) /* 每个子空间维度 */
#define K_CENT    256          /* 每个子空间的质心数 */
#define KMEANS_IT 20           /* k-means 迭代次数 */

/* ========== 类型 ========== */
typedef unsigned char pqcode_t;

typedef struct {
    float codebook[M_SUB][K_CENT][DSUB];  /* M 个码本 */
} PQCodebook;

/* ========== 工具函数 ========== */
static float l2_sqr(const float *a, const float *b, int d)
{
    float s = 0.0f;
    for (int i = 0; i < d; i++) {
        float diff = a[i] - b[i];
        s += diff * diff;
    }
    return s;
}

static int find_nearest(const float *vec, const float centroids[][DSUB],
                        int n_cent)
{
    int best = 0;
    float best_dist = FLT_MAX;
    for (int k = 0; k < n_cent; k++) {
        float d = l2_sqr(vec, centroids[k], DSUB);
        if (d < best_dist) {
            best_dist = d;
            best = k;
        }
    }
    return best;
}

/* ========== k-means 训练(单子空间) ========== */
static void kmeans_train(const float *data, int n, int d,
                         float centroids[][DSUB], int k, int max_iter)
{
    /* 随机初始化:取前 k 个向量 */
    for (int i = 0; i < k && i < n; i++)
        memcpy(centroids[i], data + i * d, d * sizeof(float));

    int *assign = (int *)malloc(n * sizeof(int));
    float *sums = (float *)calloc(k * d, sizeof(float));
    int *counts = (int *)calloc(k, sizeof(int));

    for (int iter = 0; iter < max_iter; iter++) {
        /* 分配 */
        for (int i = 0; i < n; i++)
            assign[i] = find_nearest(data + i * d, centroids, k);

        /* 重计算质心 */
        memset(sums, 0, k * d * sizeof(float));
        memset(counts, 0, k * sizeof(int));

        for (int i = 0; i < n; i++) {
            int c = assign[i];
            counts[c]++;
            for (int j = 0; j < d; j++)
                sums[c * d + j] += data[i * d + j];
        }
        for (int c = 0; c < k; c++) {
            if (counts[c] == 0) continue;
            for (int j = 0; j < d; j++)
                centroids[c][j] = sums[c * d + j] / counts[c];
        }
    }

    free(assign);
    free(sums);
    free(counts);
}

/* ========== PQ 码本训练 ========== */
void pq_train(const float *train_data, int n_train, PQCodebook *cb)
{
    float *sub_data = (float *)malloc(n_train * DSUB * sizeof(float));

    for (int m = 0; m < M_SUB; m++) {
        /* 提取第 m 个子空间的数据 */
        for (int i = 0; i < n_train; i++)
            memcpy(sub_data + i * DSUB,
                   train_data + i * DIM + m * DSUB,
                   DSUB * sizeof(float));

        /* 对该子空间运行 k-means */
        kmeans_train(sub_data, n_train, DSUB,
                     cb->codebook[m], K_CENT, KMEANS_IT);
    }

    free(sub_data);
}

/* ========== PQ 编码 ========== */
void pq_encode(const float *vec, const PQCodebook *cb, pqcode_t *code)
{
    for (int m = 0; m < M_SUB; m++) {
        const float *sub = vec + m * DSUB;
        code[m] = (pqcode_t)find_nearest(sub, cb->codebook[m], K_CENT);
    }
}

/* ========== PQ 批量编码 ========== */
void pq_encode_batch(const float *vecs, int n,
                     const PQCodebook *cb, pqcode_t *codes)
{
    for (int i = 0; i < n; i++)
        pq_encode(vecs + i * DIM, cb, codes + i * M_SUB);
}

/* ========== ADC 查找表构建 ========== */
void adc_build_lut(const float *query, const PQCodebook *cb,
                   float lut[M_SUB][K_CENT])
{
    for (int m = 0; m < M_SUB; m++) {
        const float *q_sub = query + m * DSUB;
        for (int k = 0; k < K_CENT; k++)
            lut[m][k] = l2_sqr(q_sub, cb->codebook[m][k], DSUB);
    }
}

/* ========== ADC 距离计算(单向量) ========== */
static float adc_distance(const pqcode_t *code,
                          const float lut[M_SUB][K_CENT])
{
    float dist = 0.0f;
    for (int m = 0; m < M_SUB; m++)
        dist += lut[m][code[m]];
    return dist;
}

/* ========== ADC 搜索(Top-K) ========== */
typedef struct {
    int id;
    float dist;
} Result;

static int cmp_result(const void *a, const void *b)
{
    float da = ((const Result *)a)->dist;
    float db = ((const Result *)b)->dist;
    return (da > db) - (da < db);
}

int adc_search(const float *query, const pqcode_t *codes, int n_db,
               const PQCodebook *cb, int topk, Result *results)
{
    float lut[M_SUB][K_CENT];
    adc_build_lut(query, cb, lut);

    /* 简单实现:计算所有距离后排序取 top-k */
    Result *all = (Result *)malloc(n_db * sizeof(Result));
    for (int i = 0; i < n_db; i++) {
        all[i].id = i;
        all[i].dist = adc_distance(codes + i * M_SUB, lut);
    }

    qsort(all, n_db, sizeof(Result), cmp_result);

    int ret = topk < n_db ? topk : n_db;
    memcpy(results, all, ret * sizeof(Result));
    free(all);
    return ret;
}

/* ========== 验证用的精确搜索 ========== */
static float exact_l2(const float *a, const float *b)
{
    return l2_sqr(a, b, DIM);
}

/* ========== 主函数:演示完整流程 ========== */
int main(void)
{
    srand((unsigned)time(NULL));

    const int N_TRAIN = 10000;
    const int N_DB    = 50000;
    const int N_QUERY = 5;
    const int TOPK    = 10;

    printf("PQ 编码器 + ADC 搜索演示\n");
    printf("维度=%d  子空间=%d  质心=%d\n", DIM, M_SUB, K_CENT);
    printf("训练集=%d  数据库=%d  查询=%d  Top-K=%d\n\n",
           N_TRAIN, N_DB, N_QUERY, TOPK);

    /* 生成随机数据 */
    float *train_data = (float *)malloc(N_TRAIN * DIM * sizeof(float));
    float *db_data    = (float *)malloc(N_DB * DIM * sizeof(float));
    float *queries    = (float *)malloc(N_QUERY * DIM * sizeof(float));

    for (int i = 0; i < N_TRAIN * DIM; i++)
        train_data[i] = (float)rand() / RAND_MAX - 0.5f;
    for (int i = 0; i < N_DB * DIM; i++)
        db_data[i] = (float)rand() / RAND_MAX - 0.5f;
    for (int i = 0; i < N_QUERY * DIM; i++)
        queries[i] = (float)rand() / RAND_MAX - 0.5f;

    /* 训练码本 */
    PQCodebook cb;
    printf("训练 PQ 码本...\n");
    clock_t t0 = clock();
    pq_train(train_data, N_TRAIN, &cb);
    double train_ms = (double)(clock() - t0) / CLOCKS_PER_SEC * 1000;
    printf("训练完成,耗时 %.1f ms\n\n", train_ms);

    /* 编码数据库 */
    pqcode_t *codes = (pqcode_t *)malloc(N_DB * M_SUB);
    printf("编码 %d 个向量...\n", N_DB);
    t0 = clock();
    pq_encode_batch(db_data, N_DB, &cb, codes);
    double enc_ms = (double)(clock() - t0) / CLOCKS_PER_SEC * 1000;
    printf("编码完成,耗时 %.1f ms\n", enc_ms);
    printf("原始大小: %.2f MB  PQ 大小: %.2f MB  压缩比: %.0fx\n\n",
           (double)N_DB * DIM * 4 / 1e6,
           (double)N_DB * M_SUB / 1e6,
           (double)(DIM * 4) / M_SUB);

    /* ADC 搜索并与精确搜索对比 */
    Result *results = (Result *)malloc(TOPK * sizeof(Result));

    for (int qi = 0; qi < N_QUERY; qi++) {
        const float *q = queries + qi * DIM;

        printf("--- 查询 %d ---\n", qi);

        /* ADC 搜索 */
        t0 = clock();
        adc_search(q, codes, N_DB, &cb, TOPK, results);
        double search_ms = (double)(clock() - t0) / CLOCKS_PER_SEC * 1000;

        /* 精确搜索找 top-1 */
        int exact_best = 0;
        float exact_best_dist = FLT_MAX;
        for (int i = 0; i < N_DB; i++) {
            float d = exact_l2(q, db_data + i * DIM);
            if (d < exact_best_dist) {
                exact_best_dist = d;
                exact_best = i;
            }
        }

        printf("ADC top-1: id=%d  approx_dist=%.4f  搜索耗时=%.2f ms\n",
               results[0].id, results[0].dist, search_ms);
        printf("精确 top-1: id=%d  exact_dist=%.4f\n",
               exact_best, exact_best_dist);

        /* 检查精确 top-1 是否在 ADC top-10 中 */
        int found = 0;
        for (int r = 0; r < TOPK; r++) {
            if (results[r].id == exact_best) {
                found = 1;
                printf("精确 top-1 出现在 ADC 结果的第 %d\n", r + 1);
                break;
            }
        }
        if (!found)
            printf("精确 top-1 未出现在 ADC top-%d\n", TOPK);
        printf("\n");
    }

    /* 清理 */
    free(train_data);
    free(db_data);
    free(queries);
    free(codes);
    free(results);

    printf("内存占用对比:\n");
    printf("  float32 原始向量: %d bytes/vec\n", DIM * 4);
    printf("  PQ 编码:          %d bytes/vec\n", M_SUB);
    printf("  压缩比:           %dx\n", DIM * 4 / M_SUB);

    return 0;
}

编译和运行:

gcc -O2 -o pq_demo pq_demo.c -lm
./pq_demo

这个实现大约 200 行,涵盖了 PQ 的核心流程。生产环境中还需要加入 SIMD 加速、多线程、堆排序替代全排序等优化,但核心算法逻辑是一致的。

十、基准测试:IVF-PQ vs HNSW

测试环境与数据集

使用 SIFT1B(10 亿 128 维 uint8 向量,转 float32)作为基准数据集,10000 条查询向量,ground truth 为精确暴力搜索结果。

测试机器:128 GB RAM,32 核 Xeon,单线程搜索。

结果对比

指标 HNSW(M=32, ef=128) IVF-PQ(nlist=65536, nprobe=64, M=8) IVF-PQ(nprobe=128, M=16)
索引内存 ~600 GB ~16 GB ~24 GB
构建时间 ~8 小时 ~2 小时 ~3 小时
0.99 0.62 0.78
0.98 0.85 0.93
0.97 0.95 0.98
QPS(单线程) 800 3200 1800
QPS(32 线程) 18000 85000 48000

分析

HNSW 的优势: - 精度极高,几乎等于暴力搜索 - 搜索路径短,延迟稳定

HNSW 的劣势: - 内存开销巨大,600 GB 意味着需要高配服务器 - 构建时间长,图结构无法增量更新(严格说可以,但很慢)

IVF-PQ 的优势: - 内存效率惊人,十亿向量 16 GB - 吞吐量高,特别是在多线程下 - 构建时间相对短,倒排列表易于增量更新

IVF-PQ 的劣势: - 偏低,因为 PQ 的量化误差在细粒度排序上影响大 - 需要仔细调参(nlist、nprobe、M 的组合)

什么时候用什么

数据量 < 100 万        --> HNSW 或 brute-force
数据量 1M - 100M       --> HNSW(内存够)或 IVF-PQ(内存紧张)
数据量 100M - 1B       --> IVF-PQ 是默认选择
数据量 > 1B            --> IVF-PQ + 分片,或者 DiskANN

Re-ranking 策略

IVF-PQ 的常见做法是两阶段检索:

  1. 粗筛:用 IVF-PQ 快速找到 top-1000 候选
  2. 精排:从磁盘或二级存储加载候选向量的原始数据,做精确距离计算,重新排序取 top-10

这样可以把 从 0.85 提升到 0.97 以上,同时保持高吞吐。Faiss 的 IndexRefine 就是做这个的:

# 两阶段:IVF-PQ 粗筛 + Flat 精排
index = faiss.index_factory(128, "IVF65536,PQ8,Refine(Flat)")

十一、工业实践

Facebook / Meta 的大规模相似搜索

Facebook 的视觉相似性检索系统是 IVF-PQ 最大规模的生产部署之一。他们在数十亿张图片的嵌入向量上运行 Faiss,用于:

他们的经验表明,OPQ + IVF-PQ + re-ranking 的组合在 > 0.95 的情况下,单台机器可以服务数十亿向量。

Spotify 的音乐推荐

Spotify 将每首歌映射为一个嵌入向量(结合音频特征和用户行为),然后用 IVF-PQ 在嵌入空间中做 kNN 搜索,为用户找到相似的歌曲。数亿首歌的嵌入空间,IVF-PQ 是为数不多能在合理内存下运行的方案。

工程陷阱一览

陷阱 症状 解决方案
训练数据分布与实际数据不匹配 recall 远低于预期 从实际数据中均匀采样训练集,而非用合成数据
nlist 过大,每个倒排列表太短 搜索速度快但 recall 极低 经验法则:\(\text{nlist} \approx \sqrt{N}\)\(4\sqrt{N}\)
nlist 过小,每个倒排列表太长 搜索速度慢 增大 nlist,或使用 HNSW 作为粗量化器
nprobe 设太小 recall 不达标 逐步增大 nprobe,画 recall-QPS 曲线找拐点
PQ 的 M 太小 量化误差大,recall 低 增大 M(代价是更多内存),或使用 OPQ
忘记做残差编码 白白损失 2-5 个点的 recall Faiss 默认对 IVF-PQ 做残差编码,自研系统别忘了
向量未归一化 cosine 距离结果错误 搜索前 L2 归一化,然后用内积距离等价替代
训练集太小 码本质量差 至少用 \(K \times 40\) 个样本训练 PQ,\(L \times 40\) 个样本训练粗量化器
数据倾斜导致倒排列表不平衡 某些 cell 极慢 使用多级量化或对大 cell 做二次分裂
在线更新导致索引退化 recall 随时间下降 定期重训码本和粗量化器,或使用增量训练策略
GPU 显存不足 OOM 崩溃 分批训练,或使用 CPU-GPU 混合模式
忽略预热 首次查询极慢 启动时做一批虚拟查询,预热码本和倒排列表到缓存

十二、个人思考

写完这篇文章,我想谈几点个人看法。

PQ 是信息论的胜利

从信息论的角度看,PQ 的本质是对向量空间做有损压缩。它的巧妙之处在于利用了乘积结构——用 \(M \times K\) 个参数描述 \(K^M\) 个码字,这是一种极其高效的码本设计。Shannon 的率失真理论告诉我们,对于给定的码率(每向量 \(M \log_2 K\) 比特),存在一个理论最优的量化失真下界。PQ 不一定达到这个下界,但它给出了一个在计算上可行的、工程上实用的近似。

HNSW 不会被取代,IVF-PQ 也不会

这两种方法解决的是不同规模的问题。HNSW 在千万级以下几乎是默认选择——精度高、调参简单、延迟稳定。IVF-PQ 在亿级以上几乎是唯一选择——没有它你根本装不下数据。

真正有意思的是中间地带:1000 万到 1 亿之间。这个范围内,HNSW 的内存开销还没到不可接受的地步,但 IVF-PQ 的吞吐优势已经开始显现。选择哪个取决于你更在乎精度还是成本。

未来在磁盘上

内存终究是有限的。当数据量增长到百亿甚至千亿级,即便是 IVF-PQ 的 8 字节/向量也装不下。这时候需要的是 DiskANN 这样的方案——把 PQ 码放在内存里做粗筛,把完整向量放在 SSD 上做精排。一次搜索只需要几次随机 SSD 读取,延迟可以控制在个位数毫秒。

Microsoft 的 DiskANN 和 Google 的 ScaNN 正在这个方向上推进。下一篇文章我们会详细讨论它们。

工程师的权衡艺术

最后,我想说:向量检索的核心不是某个算法有多精妙,而是工程师如何在精度、速度、内存、成本之间找到最优的权衡点。IVF-PQ 给了你一个很大的调参空间:nlist、nprobe、M、K、是否用 OPQ、是否做 re-ranking。每个参数的调整都是在某两个指标之间做取舍。

没有万能的配置,只有适合你场景的配置。理解算法原理,才能在这个多维权衡空间中找到最优解。

附录:关键公式速查

符号 含义 典型值
\(D\) 向量维度 128
\(M\) PQ 子空间数 8, 16, 32
\(K\) 每个子空间的质心数 256
\(L\)(nlist) IVF 粗量化簇数 \(\sqrt{N}\)\(4\sqrt{N}\)
nprobe 搜索时探测的簇数 \(\sqrt{L}\)\(4\sqrt{L}\)
压缩比 \(D \times 4 / M\) 64x(\(D=128, M=8\)
PQ 码本大小 \(M \times K \times (D/M) \times 4\) 128 KB
倒排列表总大小 \((M + \text{id\_size}) \times N\) 16 GB(\(N=10^9, M=8\)
ADC 查找表大小 \(M \times K \times 4\) 8 KB
有效码字数 \(K^M\) \(\approx 1.8 \times 10^{19}\)

参考文献

  1. Jegou, H., Douze, M., & Schmid, C. (2011). Product Quantization for Nearest Neighbor Search. IEEE TPAMI.
  2. Ge, T., He, K., Ke, Q., & Sun, J. (2013). Optimized Product Quantization. IEEE TPAMI.
  3. Baranchuk, D., Babenko, A., & Malkov, Y. (2018). Revisiting the Inverted Indices for Billion-Scale Approximate Nearest Neighbors. ECCV.
  4. Johnson, J., Douze, M., & Jegou, H. (2019). Billion-Scale Similarity Search with GPUs. IEEE TBD.
  5. Faiss Wiki: https://github.com/facebookresearch/faiss/wiki

上一篇: HNSW 下一篇: ScaNN 与 DiskANN 相关阅读: - HNSW - 整数压缩


By .