Embedding 量化

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

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

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

二值量化

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

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

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

Sentence Transformers 中的二值量化

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

因此，实际上将维度为 1024 的 float32 embedding 量化会产生维度为 128 的 int8 或 uint8 embedding。请参阅下面两种使用 Sentence Transformers 生成量化 embedding 的方法

from sentence_transformers import SentenceTransformer
from sentence_transformers.quantization import quantize_embeddings

# 1. Load an embedding model
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

# 2a. Encode some text using "binary" quantization
binary_embeddings = model.encode(
    ["I am driving to the lake.", "It is a beautiful day."],
    precision="binary",
)

# 2b. or, encode some text without quantization & apply quantization afterwards
embeddings = model.encode(["I am driving to the lake.", "It is a beautiful day."])
binary_embeddings = quantize_embeddings(embeddings, precision="binary")

在这里您可以看到默认 float32 embedding 和二值 embedding 在形状、大小和 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 embedding 转换为 int8，我们使用一种称为标量量化的过程。这涉及到将 float32 值的连续范围映射到 int8 值的离散集合，该集合可以表示 256 个不同的级别（从 -128 到 127）。这是通过使用一个大型的 embedding 校准数据集来完成的。我们计算这些 embedding 的范围，即每个 embedding 维度的 min 和 max。从那里，我们计算我们对每个值进行分类的步长（桶）。

为了进一步提升检索性能，您可以选择性地应用与二值 embedding 相同的重排序步骤。这里需要注意的是，校准数据集对性能有很大影响，因为它定义了桶。

Sentence Transformers 中的标量量化

将维度为 1024 的 embedding 量化为 int8 会产生 1024 字节。实际上，我们可以选择 uint8 或 int8。这个选择通常取决于您的向量库/数据库支持什么。

实际上，建议为标量量化提供以下任一选项

  1. 一大批 embedding 一次性量化，或

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

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

如果这些都不是，您将收到如下警告

Computing int8 quantization buckets based on 2 embeddings. int8 quantization is more stable with 'ranges' calculated from more embeddings or a 'calibration_embeddings' that can be used to calculate the buckets.

请参阅下面如何使用 Sentence Transformers 生成标量量化 embedding

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

# 1. Load an embedding model
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

# 2. Prepare an example calibration dataset
corpus = load_dataset("nq_open", split="train[:1000]")["question"]
calibration_embeddings = model.encode(corpus)

# 3. Encode some text without quantization & apply quantization afterwards
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 embedding 和 int8 标量 embedding 在形状、大小和 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

结合二值量化和标量量化

可以将二值量化和标量量化结合起来，以获得两全其美的效果：二值 embedding 的极快速度和标量 embedding 通过重排序实现的出色性能保留。请参阅下面的演示，了解这种方法涉及维基百科 4100 万文本的实际实现。该设置的管道如下：

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

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

  3. 使用量化查询在二值索引（41M 二值 embedding；5.2GB 内存/磁盘空间）中搜索前 40 个文档。

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

  5. 使用 float32 查询和 int8 embedding 对前 40 个文档进行重排序，以获得前 10 个文档。

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

通过这种方法，我们为索引使用了 5.2GB 内存和 52GB 磁盘空间。这比常规检索少得多，常规检索需要 200GB 内存和 200GB 磁盘空间。特别是当您进一步扩展时，这将显著降低延迟和成本。

其他扩展

请注意，embedding 量化可以与其他方法结合使用以提高检索效率，例如俄罗斯套娃 Embedding。此外，检索与重排序也与量化 embedding 配合得很好，即您仍然可以使用 Cross-Encoder 进行重排序。

演示

以下演示展示了通过结合二值搜索和标量（int8）重排序的 exact 搜索的检索效率。该解决方案需要 5GB 内存用于二值索引，50GB 磁盘空间用于二值和标量索引，这比常规 float32 检索所需的 200GB 内存和磁盘空间少得多。此外，检索速度也快得多。

自己尝试

以下脚本可用于实验 embedding 量化以进行检索及其他。共有三类