乘积量化 (PQ) 工作原理

存储和查询数百万乃至数十亿高维向量 (vector)，例如由大型语言模型（LLM）嵌入 (embedding)模型生成的向量，带来很大的问题。这些向量庞大的内存占用（通常是float32类型，每个向量需要 $D \times 4$ 字节）可能令人望而却步，而且计算查询向量与所有数据库向量之间的精确距离在计算上成本很高。乘积量化 (quantization)（PQ）是一种向量压缩技术，旨在处理这两个问题，能够大幅节省内存并加快近似距离计算。

PQ 遵循“分而治之”的原则。它不是一次性量化整个高维向量（这通常会导致分辨率不佳），而是首先将向量分解成多个低维子向量。

向量 (vector)分解

假设有一个 $D$ 维空间中的原始向量 $x$ ($x \in \mathbb{R}^D$)。PQ 将这个向量分成 $M$ 个不同的子向量，每个子向量的维度为 $D/M$。为简单起见，我们假设 $D$ 可以被 $M$ 整除。

$$x = [x_1 | x_2 | \dots | x_M]$$

这里，每个 $x_m$ ($1 \le m \le M$) 都是一个维度为 $D/M$ 的子向量。比如，一个 768 维的向量可以分成 $M=48$ 个子向量，每个维度为 $D/M = 16$。

子空间量化 (quantization)和码本

核心思想是独立地量化这 $M$ 个子空间中的每一个。对于每个子空间 $m$，我们创建一个独立的码本 $C_m$。这个码本包含 $k^*$ 个代表向量 (vector)（通常称为聚类中心），它们是该子空间特有的。这些聚类中心通常是使用 k-means 聚类算法，从与该子空间对应的子向量训练集中学习得到的。

$$C_m = \{c_{m,1}, c_{m,2}, \dots, c_{m,k^*}\} \quad \text{其中 } c_{m,j} \in \mathbb{R}^{D/M}$$

聚类中心的数量 $k^*$ 是一个重要的参数 (parameter)。它决定了每个子空间内量化的粒度。一个常见的选择是 $k^*=256$，这使得每个聚类中心索引可以用一个字节（8 位）表示。

生成 $M$ 个码本后，任何原始向量 $x$ 都可以被编码。向量 $x$ 的每个子向量 $x_m$ 被映射到其在对应码本 $C_m$ 中最近聚类中心的索引（ID）。令 $q_m(x_m)$ 为查找子向量 $x_m$ 在 $C_m$ 中最近聚类中心索引的函数：

$$q_m(x_m) = \underset{j \in \{1, \dots, k^*\}}{\operatorname{argmin}} \| x_m - c_{m,j} \|^2$$

原始向量 $x$ 随后由 $M$ 个聚类中心索引序列表示：

$$\text{编码}(x) = [q_1(x_1), q_2(x_2), \dots, q_M(x_M)]$$

每个索引 $q_m(x_m)$ 需要 $\log_2(k^*)$ 位。如果我们选择 $k^*=256$，每个索引需要 8 位（1 字节）。因此，整个 $D$ 维向量 $x$ 被压缩成仅仅 $M$ 字节。与存储原始 $D$ 维 float32 向量（需要 $D \times 4$ 字节）相比，这代表了很大的压缩比，通常根据 $D$ 和 $M$ 的不同，范围在 10 倍到 100 倍之间。对于我们 768 维的例子，当 $M=48$ 时，压缩后的大小是 48 字节，而原始大小为 $768 \times 4 = 3072$ 字节，压缩因子为 64 倍。

乘积量化的过程：一个高维向量被分割成子向量。每个子空间进行 k-means 聚类以创建码本。原始向量随后由每个子空间中最近聚类中心 ID 的序列表示。

近似距离计算

PQ 的主要优点，除了内存压缩之外，是它能够快速近似计算查询向量 (vector) $q$ 与压缩数据库向量 $x$（由其编码 $\text{编码}(x)$ 表示）之间的距离。

非对称距离计算 (ADC)

最常用的方法是非对称距离计算（ADC）。在这种方法中，查询向量 $q$ 保持其原始未压缩的形式。到数据库向量 $x$（由其 PQ 编码表示）的距离是通过求和查询的子向量 $q_m$ 与存储的 $x$ 编码对应的聚类中心之间的距离来近似计算的。

具体来说，令 $\text{编码}(x) = [id_1, id_2, \dots, id_M]$，其中 $id_m = q_m(x_m)$。近似欧几里得距离平方计算如下：

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

这里的主要优化是，对于所有 $m \in \{1, \dots, M\}$ 和 $j \in \{1, \dots, k^*\}$，距离 $\| q_m - c_{m,j} \|^2$ 可以对每个查询 $q$ 只预先计算一次。这会创建 $M$ 个查找表，每个表包含 $k^*$ 个预计算距离。然后，数据库向量 $x$ 的总距离近似值就简化为 $M$ 次表查找和 $M-1$ 次加法。这比计算完整的 $D$ 维距离快很多。

对称距离计算 (SDC)

另一种选择是对称距离计算（SDC），其中查询向量 $q$ 也使用相同的码本进行量化 (quantization)。然后仅使用聚类中心来近似计算距离：

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

SDC 可以更快，因为它完全在压缩域中操作，可能会使用每个子空间内所有聚类中心对之间的预计算距离。然而，量化查询会引入额外的近似误差，通常使得 SDC 不如 ADC 准确。当优先考虑准确性时，通常首选 ADC。

权衡与考量

乘积量化 (quantization)提供了很大的好处，但也伴随着固有的权衡：

1. 准确性与压缩： PQ 是一种有损压缩技术。将子向量 (vector)映射到聚类中心时会丢失信息。这种损失直接影响最近邻搜索的准确性（召回率）。控制这种权衡的主要因素是：

   • M（子空间数量）： 增加 $M$（减少子空间维度 $D/M$）通常会提高准确性，但会降低压缩比（因为编码长度是 $M$ 字节，假设 $k^*=256$）。
   • k*（码本大小）： 增加 $k^*$ 可以在每个子空间内实现更细粒度的量化，从而提高准确性，但会增加码本所需的内存以及每个编码索引所需的位长。$k^*=256$ 是一个常见的实用选择，它在准确性和字节对齐 (alignment)之间取得了平衡。

2. 训练成本： 生成 $M$ 个码本需要对从代表性训练样本中提取的、可能很大的子向量数据集运行 k-means 聚类 $M$ 次。这个离线训练步骤可能计算量很大。

3. 与索引结构结合： 尽管 PQ 提供了压缩和快速近似距离计算，但它本身不提供像 HNSW 或 IVF 那样的亚线性搜索结构。PQ 通常与这些结构结合使用。例如，在 IVFPQ（下一节将介绍）中，PQ 用于压缩存储在 IVF 索引倒排列表中的向量，从而大幅减少索引大小并在搜索细化阶段加快距离计算。类似地，PQ 可以压缩存储在 HNSW 图节点上的向量。

乘积量化是管理现代向量数据集规模的一种重要技术。通过分解向量并独立量化子空间，它实现了令人称赞的压缩效果，并能够快速进行近似距离计算，使其成为构建大规模、高效向量搜索系统中不可或缺的工具。理解其工作原理对于优化要求高的 LLM 应用中的内存使用和查询延迟很重要。