嵌入量化

嵌入可能难以扩展，这导致昂贵的解决方案和高延迟。目前，许多最先进的模型生成具有 1024 维的嵌入，每个维度都编码为 float32，即每个维度需要 4 字节。要在 5000 万个向量上进行检索，因此需要大约 200GB 的内存。这通常需要复杂且成本高昂的大规模解决方案。

然而，有一种新方法可以应对这个问题；它涉及减少嵌入中每个单独值的大小：量化。量化实验表明，我们可以在显著加快计算速度并节省内存、存储和成本的同时，保持大量性能。

要了解更多关于嵌入量化及其性能的信息，请阅读 Sentence Transformers 和 mixedbread.ai 的博客文章。

二值量化

二值量化指的是将嵌入中的 float32 值转换为 1 位值，从而使内存和存储使用量减少了 32 倍。要将 float32 嵌入量化为二值，我们只需在 0 处对归一化嵌入进行阈值处理：如果值大于 0，我们将其设为 1，否则将其转换为 0。我们可以使用汉明距离（Hamming Distance）来高效地执行这些二值嵌入的检索。汉明距离简单来说就是两个二值嵌入在多少个位置上的比特不同。汉明距离越低，嵌入越接近，因此文档越相关。汉明距离的一个巨大优势是它可以用 2 个 CPU 周期轻松计算，从而实现极快的性能。

Yamada 等人（2021）引入了一个重新评分步骤，他们称之为重新排序，以提升性能。他们提出 float32 查询嵌入可以与二值文档嵌入使用点积进行比较。实际上，我们首先使用二值查询嵌入和二值文档嵌入检索 rescore_multiplier * top_k 结果——即双重二值检索的前 k 个结果列表——然后使用 float32 查询嵌入对这些二值文档嵌入列表进行重新评分。

通过应用这一新颖的重新评分步骤，我们能够保留高达 ~96% 的总检索性能，同时将内存和磁盘空间使用量减少 32 倍，并将检索速度提高高达 32 倍。

Sentence Transformers 中的二值量化

将 1024 维度的嵌入量化为二值将产生 1024 位。实际上，将比特存储为字节更为常见，因此当我们量化为二值嵌入时，我们使用 np.packbits 将比特打包为字节。

因此，实际上，将 1024 维度的 float32 嵌入量化将产生一个 128 维度的 int8 或 uint8 嵌入。请参阅以下两种方法，了解如何使用 Sentence Transformers 生成量化嵌入：

from sentence_transformers import SentenceTransformer
from sentence_transformers.quantization import quantize_embeddings

# 1. 加载嵌入模型
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

# 2a. 使用 "binary" 量化编码一些文本
binary_embeddings = model.encode(
    ["I am driving to the lake.", "It is a beautiful day."],
    precision="binary",
)

# 2b. 或者，先编码一些文本而不进行量化，然后应用量化
embeddings = model.encode(["I am driving to the lake.", "It is a beautiful day."])
binary_embeddings = quantize_embeddings(embeddings, precision="binary")

在这里，您可以看到默认的 float32 嵌入与二值嵌入在形状、大小和 numpy dtype 方面的差异：

>>> embeddings.shape
(2, 1024)
>>> embeddings.nbytes
8192
>>> embeddings.dtype
float32
>>> binary_embeddings.shape
(2, 128)
>>> binary_embeddings.nbytes
256
>>> binary_embeddings.dtype
int8

请注意，您也可以选择 "ubinary" 以使用无符号 uint8 数据格式进行量化为二值。这可能是您的向量库/数据库的要求。

标量（int8）量化

要将 float32 嵌入转换为 int8，我们使用一种称为标量量化的过程。这涉及将 float32 值的连续范围映射到 int8 值的离散集合，后者可以表示 256 个不同的级别（从 -128 到 127）。这是通过使用一个包含大量嵌入的校准数据集来完成的。我们计算这些嵌入的范围，即每个嵌入维度的 min 和 max。从那里，我们计算分类每个值的步长（桶）。

为了进一步提高检索性能，您可以选择性地应用与二进制嵌入相同的重新评分步骤。需要注意的是，校准数据集对性能有很大影响，因为它定义了桶。

Sentence Transformers 中的标量量化

将维度为 1024 的嵌入量化为 int8 后，结果为 1024 字节。在实践中，我们可以选择 uint8 或 int8。这一选择通常取决于您的向量库/数据库支持的内容。

在实践中，建议使用以下任一方式提供标量量化：

  1. 一个大型嵌入集，一次性量化所有嵌入，或

  2. 每个嵌入维度的 min 和 max 范围，或

  3. 一个大型校准数据集，从中可以计算 min 和 max 范围。

如果这些情况都不适用，您将收到类似以下的警告：

基于 2 个嵌入计算 int8 量化桶。int8 量化在从更多嵌入计算 'ranges' 或 'calibration_embeddings' 可以用于计算桶时更加稳定。

请参见以下如何使用 Sentence Transformers 生成标量量化嵌入的示例：

from sentence_transformers import SentenceTransformer
from sentence_transformers.quantization import quantize_embeddings
from datasets import load_dataset

# 1. 加载嵌入模型
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

# 2. 准备一个示例校准数据集
corpus = load_dataset("nq_open", split="train[:1000]")["question"]
calibration_embeddings = model.encode(corpus)

# 3. 编码一些文本而不进行量化，然后应用量化
embeddings = model.encode(["I am driving to the lake.", "It is a beautiful day."])
int8_embeddings = quantize_embeddings(
    embeddings,
    precision="int8",
    calibration_embeddings=calibration_embeddings,
)

在这里，您可以看到默认的 float32 嵌入与 int8 标量嵌入在形状、大小和 numpy dtype 方面的差异：

>>> embeddings.shape
(2, 1024)
>>> embeddings.nbytes
8192
>>> embeddings.dtype
float32
>>> int8_embeddings.shape
(2, 1024)
>>> int8_embeddings.nbytes
2048
>>> int8_embeddings.dtype
int8

结合二进制和标量量化

可以将二进制和标量量化结合起来，以获得两者的最佳效果：二进制嵌入的极端速度和标量嵌入在重新评分时的高性能保留。请参见下面的演示，了解涉及 4100 万条来自 Wikipedia 的文本的实际实现。该设置的流程如下：

  1. 使用 mixedbread-ai/mxbai-embed-large-v1 SentenceTransformer 模型嵌入查询。

  2. 使用 sentence-transformers 库中的 quantize_embeddings 函数将查询量化为二进制。

  3. 使用量化的查询搜索二进制索引（41M 二进制嵌入；5.2GB 内存/磁盘空间）以获取前 40 个文档。

  4. 从磁盘上的 int8 索引（41M int8 嵌入；0 字节内存，47.5GB 磁盘空间）实时加载前 40 个文档。

  5. 使用 float32 查询和 int8 嵌入对前 40 个文档进行重新评分，以获取前 10 个文档。

  6. 按分数对前 10 个文档进行排序并显示。

通过这种方法，我们使用了 5.2GB 内存和 52GB 磁盘空间用于索引。这比正常检索所需的 200GB 内存和 200GB 磁盘空间要少得多。特别是随着规模的进一步扩大，这将显著降低延迟和成本。

其他扩展

需要注意的是，嵌入量化可以与其他方法结合以提高检索效率，例如套娃嵌入。此外，检索与重排序方法也与量化嵌入配合得非常好，即你仍然可以使用交叉编码器进行重排序。

演示

以下演示展示了通过结合二分搜索与标量（int8）重评分来使用精确搜索的检索效率。该方案需要5GB的内存用于二分索引，50GB的磁盘空间用于二分和标量索引，远低于常规float32检索所需的200GB内存和磁盘空间。此外，检索速度也快得多。

自行尝试

以下脚本可用于试验嵌入量化在检索及其他方面的应用。分为三类：