非对称量化：实现近无损检索且存储降低97%

# 非对称量化：以 97% 的存储压缩实现近乎无损的延迟交互检索 来源：https://www.mixedbread.com/blog/asymmetric-quant 像 Wholembed v3 (https://www.mixedbread.com/blog/wholembed-v3) 这样的延迟交互模型使得检索更加精确，因为它们保留了细粒度信息，而不是将整个文档压缩成一个向量。但它们也改变了存储的经济性。一个文档会产生不止一个嵌入，根据文档的复杂程度，它可能产生数百或数千个向量。每个向量都必须存储并用于后续检索。 Mixedbread Search 运行在 silo (https://www.mixedbread.com/blog/multimodal-late-interaction-billion-scale#silo-an-s3-native-retrieval-engine) 上，这是我们用于十亿级多模态延迟交互的检索引擎。Silo 将超过 25 亿个文档的向量存储在对象存储中，并在查询需要时将其加载到更快的层级中。在这样的规模下，每个文档的额外字节都会被重复数十亿次，并直接体现在每个文档的存储成本、分片冷启动时间以及每个查询需要读取的字节数上。我们需要在保持延迟交互价值所需的质量的同时，找到降低整个系统成本的方案。 本文将介绍非对称量化。这是使延迟交互在生产环境中实用化的优化之一。我们将查询向量保持较高精度，而将文档向量存储为二进制的符号位。在我们的内部基准测试套件中，这种方式将原始文档向量存储平均减少了 32 倍，从每个文档 393 KiB 降至 12.28 KiB，同时检索质量保持在 89.65 NDCG@10（对比 fp32 的 90.26）。 多向量文档的每个文档存储量从 fp32 的 393 KiB 降至 int8 查询对二进制文档的 12.28 KiB，减少了 32 倍（97%）。fp32 被绘制为 32 个单元格，每个单元格的大小等于 int8 对二进制单元格的大小。量化是指用更低精度的值（如 int8 甚至 1 位符号）来表示高精度浮点向量。目标是在减少数据量的同时保持排序质量。这对于 silo 尤其重要。对象存储提供了持久且低成本的持久化。为了使其适合实际工作负载，我们需要紧凑的索引以足够快地提供服务。而在文档端，数据量大小是成本的主导因素。 朴素的延迟交互成本高昂，因为它存储了更多向量。一个标准单向量嵌入（3072 维，fp32）每个文档占用 12 KiB。一个多向量表示（786 个向量，128 维）携带更多信息，但未压缩时大约大 33 倍。 | 表示方式 | 维度 | 每个文档存储 | 与 3072 维 fp32 单向量之比 | | :--- | :--- | :--- | :--- | | 单向量，fp32 | 3072 | 12,288 B / 12 KiB | 1.0x | | 单向量，int8 | 3072 | 3,072 B / 3 KiB | 0.25x | | 多向量，fp32 | 786 × 128 | 402,432 B / 393 KiB | 32.75x | | 多向量，int8 | 786 × 128 | 100,608 B / 98.25 KiB | 8.19x | | 多向量，binary | 786 × 128 | 12,576 B / 12.28 KiB | 1.02x | 这里的存储数字仅指原始向量数据。生产索引还包括文档 ID、元数据和布局开销。 使用二进制文档向量时，一个 786 token 的多向量文档仅比 3072 维 fp32 单向量大约 2%。这意味着，你可以付出大致相当于单一向量存储的成本，获得延迟交互的质量。这帮助我们改变了权衡。延迟交互变得默认即可实用，而不是仅用于那些值得承担存储成本的场景。 这对延迟交互来说并不是新方向。ColBERTv2 (https://arxiv.org/abs/2112.01488) 表明激进压缩可以在保持质量的同时减少延迟交互模型的存储占用。PLAID (https://arxiv.org/abs/2205.09707) 表明，通过优化检索和剪枝，可以将延迟交互检索的延迟降低到实用水平。对于生产系统，这两点都很重要：模型必须精确，表示必须足够便宜以在硬件中移动。 压缩文档向量可以节省整个语料库的存储、IO、缓存空间和冷启动时间。压缩查询向量几乎节省不了什么，因为查询很小、生命周期短暂且从未存储在索引中。 这也是为什么我们不将两侧都二值化。完全二值化检索是最紧凑的选择，但将查询降至单个比特会丢弃排序所依赖的幅度信息，其质量损失远大于仅二值化文档（如下文所示）。 因此，我们将查询保持在 int8 精度，仅将文档向量存储为二进制符号位。查询保持足够精确以维持排序，而文档端获得了对服务至关重要的存储缩减。 二进制文档向量更小，因此存储成本更低。 对于 int8 × int8 的评分，现代 ARM CPU 通过 NEON 点积指令 (https://developer.arm.com/Architectures/Neon) 提供直接支持。我们的 AArch64 内核使用 SDOT (https://developer.arm.com/documentation/ddi0596/2020-12/SIMD-FP-Instructions/SDOT--vector---Dot-Product-signed-arithmetic--vector-) 将 16 个 int8 乘法累加为 int32 通道，然后使用 `vaddvq_s32` 进行水平归约。对于 int8 × binary 评分，有用的恒等式更简单。如果每个文档维度存储为一个符号位，`b_i` 取值 `{-1, +1}`，那么： ```math \begin{aligned} \mathbf{q} \cdot \mathbf{b} &= \sum_i q_i b_i \\ &= 2 \sum_{i\,:\,b_i = +1} q_i - \sum_i q_i \end{aligned} ``` 因此，评分不需要为每个维度进行完整的乘法。我们只需要查询值中被正文档位选中的部分之和，以及查询的总和。 一个浮点文档向量被缩减为其符号位并打包成字节。用 int8 查询对二进制文档进行评分不需要每维乘法：它是两倍的正文档位选中的查询值之和减去查询总和。在内核中，文档符号位被打包成比特。对于 128 维路径，我们预先计算查询总和和八个查询位平面。每个文档 token 作为 16 个打包字节加载，通过 NEON 整数运算进行移位和掩码得到八个 0/1 掩码，然后使用 SDOT 对查询平面进行评分。最终分数使用上述恒等式：`2 * selected_query_sum - query_sum`。 二进制 × 二进制在计算上更便宜，因为它可以使用汉明距离，但质量损失太大，不适合我们的主要检索路径。 我们在内部检索基准套件中评估了几种精度组合。分数是套件平均的 NDCG@10，缩放为 0-100。NDCG@10（归一化折损累积增益，排名 10）衡量前 10 个结果相对于理想排序的排序质量，相关文档排名越高奖励越大，100 分表示完美排序。全精度基线平均得分为 90.26。int8 查询对二进制文档的平均得分为 89.65，下降了 0.61 分，同时文档向量存储减少了 32 倍。**性能下降极小，部分原因是 Wholembed v3 在训练时就考虑了 silo 的权衡，因此对量化具有鲁棒性**。 在运行时方面，我们在 ARM 机器上对评分内核进行了基准测试，使用 33 × 128 的查询对 1000 个文档（每个 786 × 128）的列表进行评分。表格报告了 2 次预热后 9 次测量中的中位延迟，以及相对于 fp32 基线的加速比。 | 查询格式 | 文档格式 | 平均 NDCG@10 | 与基线差值 | 每个文档存储 | 中位延迟 | 相对于 fp32 加速比 | | :--- | :--- | :--- | :--- | :--- | :--- | :--- | | Float | Float | 90.26 | – | 393 KiB | 14.20 ms | 1.00x | | Int8 | Int8 | 90.27 | +0.01 | 98.25 KiB | 4.44 ms | 3.20x | | Int8 | Binary | 89.65 | -0.61 | 12.28 KiB | 3.71 ms | 3.82x | | Binary | Binary | 83.06 | -7.20 | 12.28 KiB | 3.32 ms | 4.27x | NDCG@10 与每个文档存储量的关系图，展示了四种精度组合，x 轴为反向对数刻度（存储从左向右减少）。一条线连接 fp32（393 KiB, 90.26）、int8×int8（98.25 KiB, 90.27）和 int8×binary（12.28 KiB, 89.65）：质量几乎持平，而存储减少了 32 倍。在相同 12.28 KiB 下，binary×binary（83.06）比这条线低 7.2 分，因为将查询二值化丢弃了太多信号。 有两个有用的操作点。 如果我们想要在降低带宽和更快整数评分的同时获得最高质量，那么 int8 × int8 在此设置中几乎是无损的。它在测量噪声范围内略高于 fp32 基线，同时将文档存储减少 4 倍，在此基准测试中运行速度快 3.2 倍。 如果我们想要最佳的存储经济性，int8 × binary 是更有趣的点。它保持了大部分排序质量，同时将文档向量缩减 32 倍，运行速度比 fp32 快 3.8 倍。对于基于对象存储的系统，这直接减少了语料库端的字节量。 Binary × binary 在纸面上看起来很有吸引力。它使用与 int8 × binary 相同的 12.28 KiB 文档存储，并且速度最快（4.3 倍）。但相对于基线下降了 7.20 分，是 int8 × binary 下降幅度的十倍以上，尽管读取了完全相同的文档字节。唯一改变的是查询。实际上，它移除了太多查询信号。 非对称量化之所以有效，是因为检索系统支付查询和文档精度的方式不同。文档向量主导了系统的长期成本：它们被存储、复制、缓存、驱逐、重新加载并反复评分。查询向量是短命的，因此在查询上多花费几个比特，在每个存储的文档上节省比特，是正确的权衡。 对于 silo，这使得延迟交互检索更容易在大型规模上提供服务。更低的每个文档存储成本，更快的分片冷启动，以及一个将更多时间用于评分（从而允许更高 QPS）并减少字节移动的硬件路径。这使我们能够获得多向量表示的质量，而无需将每个文档视为一个大的 fp32 对象。 如果你对这类系统问题感兴趣，我们正在招聘 (https://mixedbread.com/careers)。