Lucene 中的标量量化：如何优化存储和搜索向量

Understanding Scalar Quantization in Lucene

自动字节量化在 Lucene 中的应用

HNSW 是一种功能强大且灵活的存储和搜索向量的方法，但它需要大量内存才能快速运行。例如，查询 100 万个 768 维度的 float32 向量大约需要 3GB 的 RAM。一旦开始搜索大量向量，这将变得非常昂贵。通过字节量化可以节省大约 75% 的内存。Lucene 以及 Elasticsearch 早已支持字节向量的索引构建，但这些向量的构建一直是用户的责任。这种情况即将改变，因为我们在 Lucene 中引入了 int8 标量量化。

标量量化基础知识

所有量化技术都被视为对原始数据的有损转换，这意味着为了节省空间会丢失一些信息。有关标量量化的详细解释，请参阅：标量量化入门。简而言之，标量量化是一种有损压缩技术，通过一些简单的数学计算可以在对召回率影响很小的情况下显著节省空间。

架构探索

习惯使用 Elasticsearch 的用户可能已经熟悉这些概念，但这里是关于搜索文档分布的快速概述。

每个 Elasticsearch 索引由 多个分片 组成。虽然每个分片只能分配给一个节点，但每个索引的多个分片可以在多个节点上实现计算并行。

每个分片是由单个 Lucene 索引 组成。一个 Lucene 索引由多个只读段组成。在索引过程中，文档会被缓冲并定期刷新到只读段中。当满足某些条件时，这些段可以在后台合并为更大的段。所有这些都是可配置的，并且有其自身的复杂性。当我们谈论段和合并时，我们指的是只读 Lucene 段和这些段的自动定期合并。深入了解 段合并和设计决策。

Lucene 中的分段量化

每个 Lucene 段存储以下内容：单个向量、HNSW 图索引、量化向量和计算的分位数。为了简洁，我们将重点介绍 Lucene 如何存储量化和原始向量。对于每个段，我们跟踪 vec 文件中的原始向量、veq 文件中的量化向量和单个修正乘数浮点数，以及 vemq 文件中关于量化的元数据。

图 1：原始向量存储文件的简化布局。

占用 dimension∗4∗numVectors 的磁盘空间，因为 float 值是 4 字节。由于我们正在量化，因此这些在 HNSW 搜索期间不会加载。它们仅在特定请求时使用（例如通过重排序进行暴力二次搜索），或在段合并期间重新量化。

图 2：.veq 文件的简化布局。

占用 (dimension+4)∗numVectors 的空间，并将在搜索期间加载到内存中。+4 字节是为了修正乘数浮点数，用于调整评分以提高准确性和召回率。

图 3：元数据文件的简化布局。

这里跟踪量化和向量配置以及该段的计算分位数。

因此，对于每个段，我们不仅存储量化向量，还存储用于生成这些量化向量的分位数和原始向量。那么，为什么我们还要保留原始向量呢？

适应变化的量化

由于 Lucene 定期刷新到只读段，每个段仅对所有数据有部分视图。这意味着计算的分位数仅适用于整个数据的样本集。如果您的样本足够代表整个语料库，这不是问题。但 Lucene 允许以各种方式对索引进行排序。因此，您可能会按某种方式排序数据，从而对每段的分位数计算产生偏差。另外，您可以随时刷新数据！您的样本集可能非常小，甚至只有一个向量。另一个复杂因素是，您可以控制何时进行合并。虽然 Elasticsearch 有配置默认和定期合并，但您可以通过 _force_merge API 随时请求合并。那么，我们如何在保持所有这些灵活性的同时，提供良好的量化效果？

Lucene 的向量量化会自动随着时间调整。由于 Lucene 设计了只读段架构，我们保证每个段中的数据未改变，并且代码中有明确的标记何时可以更新。这意味着在段合并期间，我们可以根据需要调整分位数，并可能重新量化向量。

图 4：三个具有不同分位数的示例段。

重新量化是否昂贵？确实有一些开销，但 Lucene 会智能地处理分位数，只有在必要时才会完全重新量化。以图 4 中的段为例。假设段 A 和 B 各有 1000 个文档，段 C 只有 100 个文档。Lucene 将取分位数的加权平均值，如果合并后的分位数接近段的原始分位数，我们不需要重新量化该段，可以利用新合并的分位数。

图 5：段 A 和 B 各有 1000 个文档，段 C 只有 100 个文档的合并分位数示例。

在图 5 中，我们可以看到合并后的分位数与段 A 和 B 的原始分位数非常相似。因此，不需要重新量化这些段的向量。而段 C 的分位数偏差太大，因此需要使用新合并的分位数重新量化。

在极端情况下，合并后的分位数与任何原始分位数差异显著。在这种情况下，我们将从每个段中抽取样本并重新计算分位数。

量化性能与数据

那么，它快吗？召回率还好吗？以下数据是在 GCP 的 c3-standard-8 实例上运行实验得出的。为了与 float32 进行公平比较，我们使用了足够大的实例来容纳内存中的原始向量。我们索引了 400,000 个 Cohere Wiki 向量 使用 maximum-inner-product。

图 6：量化向量与原始向量的 Recall@10。

量化向量的搜索性能明显快于原始向量，召回率通过多收集 5 个向量就可以快速恢复；见 quantized@15。

图 6 讲述了这个故事。虽然召回率有所不同，但这是预期的，并不显著。通过多收集 5 个向量，召回率差异消失。这一切都是通过 2 倍更快的段合并和 float32 向量的 1/4 内存实现的。

结论

Lucene 提供了一个独特的解决方案来解决一个困难的问题。量化不需要“训练”或“优化”步骤。在 Lucene 中，它将自动工作。不需要担心数据变化时需要“重新训练”向量索引。Lucene 会检测到显著变化，并在数据的整个生命周期内自动处理这一切。现在我们将这一功能带入 Elasticsearch！有兴趣的同学，快来尝试吧！