嵌入量化

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

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

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

二值量化

二值量化是指将嵌入中的 float32 值转换为 1 位值，从而将内存和存储使用量减少 32 倍。要将 float32 嵌入量化为二进制，我们只需将归一化嵌入阈值设置为 0：如果值大于 0，我们将其设为 1，否则我们将其转换为 0。我们可以使用汉明距离来有效地使用这些二值嵌入执行检索。这只是两个二值嵌入的位不同的位置数。汉明距离越小，嵌入越接近，因此文档越相关。汉明距离的一个巨大优势是它可以轻松地用 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. 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 万篇文本。该设置的 Pipeline 如下

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

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

  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 内存和磁盘空间。此外，检索速度更快。

亲自尝试

以下脚本可用于实验嵌入量化以进行检索及其他操作。分为三个类别