嵌入量化

嵌入的规模化可能具有挑战性，这会导致昂贵的解决方案和高延迟。目前，许多最先进的模型生成的嵌入具有 1024 个维度，每个维度都使用 float32 进行编码，即每个维度需要 4 个字节。因此，要对 5000 万个向量进行检索，您将需要大约 200GB 的内存。这往往需要大规模且复杂的解决方案。

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

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

二值量化

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

Yamada 等人 (2021) 引入了一个重打分步骤，他们称之为*重新排序 (rerank)*，以提升性能。他们提出，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. 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 嵌入和二值嵌入在形状、大小和 numpy 数据类型方面的差异。

>>> 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 范围。

如果以上情况均不满足，您将收到如下警告：

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 生成标量量化嵌入的方法：

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 嵌入和 int8 标量嵌入在形状、大小和 numpy 数据类型方面的差异。

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

结合二值和标量量化

可以结合二值和标量量化来两全其美：利用二值嵌入的极快速度和标量嵌入重打分后出色的性能保持。请参阅下面的演示，了解这种方法在涉及 4100 万条维基百科文本的实际实现。该设置的流程如下：

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

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

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

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

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

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

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

附加扩展

请注意，嵌入量化可以与其他方法结合以提高检索效率，例如套娃嵌入（Matryoshka Embeddings）。此外，检索与重排（Retrieve & Re-Rank）也与量化嵌入配合得很好，即您仍然可以使用 Cross-Encoder 进行重排。

演示

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

亲自尝试

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