嵌入量化

嵌入可能难以扩展，导致昂贵的解决方案和高延迟。目前，许多最先进的模型生成 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 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 范围。

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

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 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 万维基百科文本的这种方法的实际实现。该设置的管道如下：

  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 磁盘空间）要少得多。特别是当您进一步扩展时，这将显著降低延迟和成本。

额外扩展

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

演示

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

亲自尝试

以下脚本可用于实验嵌入量化，以实现检索及其他功能。共有三类