嵌入量化

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

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

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

二值量化

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

Yamada et al. (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。从那里，我们计算对每个值进行分类的步长（buckets）。

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

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

结合二值量化和标量量化

可以将二值量化和标量量化结合起来，以获得两全其美的效果：来自二值嵌入的极速和通过重评分对标量嵌入的卓越性能保留。请参阅下面的演示，了解这种方法的实际应用，其中涉及来自 Wikipedia 的 4100 万篇文本。该设置的 pipeline 如下

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

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

  3. 使用量化的查询搜索二值索引（4100 万个二值嵌入；5.2GB 内存/磁盘空间），以获得前 40 个文档。

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

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

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

通过这种方法，我们对索引使用 5.2GB 内存和 52GB 磁盘空间。这远小于正常检索，正常检索需要 200GB 内存和 200GB 磁盘空间。特别是当您进一步扩展时，这将显着减少延迟和成本。

其他扩展

请注意，嵌入量化可以与其他方法结合使用，以提高检索效率，例如 Matryoshka 嵌入。此外，检索 & 重排序 也非常适用于量化嵌入，即您仍然可以使用 Cross-Encoder 进行重排序。

演示

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

亲自尝试

以下脚本可用于试验嵌入量化，以进行检索及其他应用。共有三个类别