二值和标量嵌入量化：实现更快、更经济的检索

我们引入嵌入量化的概念，并展示其对检索速度、内存使用、磁盘空间和成本的影响。我们将讨论嵌入在理论上和实践中如何进行量化，然后介绍一个演示，该演示展示了对 4100 万维基百科文本的真实检索场景。

目录

• 为什么需要嵌入？
  • 嵌入可能难以扩展
• 提高可扩展性
  • 二值量化
    • Sentence Transformers 中的二值量化
    • 向量数据库中的二值量化
  • 标量 (int8) 量化
    • Sentence Transformers 中的标量量化
    • 向量数据库中的标量量化
  • 结合二值和标量量化
  • 量化实验
  • 重排序（Rescoring）的影响
    • 二值重排序
    • 标量 (Int8) 重排序
    • 检索速度
  • 性能总结
  • 演示
  • 亲自尝试
  • 未来工作
  • 致谢
  • 引用
  • 参考文献

为什么需要嵌入？

嵌入是自然语言处理中功能最强大的工具之一，支持各种设置和用例。本质上，嵌入是更复杂对象（如文本、图像、音频等）的数值表示。具体来说，这些对象被表示为 n 维向量。

转换复杂对象后，您可以通过计算相应嵌入的相似度来确定它们的相似性！这对于许多用例至关重要：它是推荐系统、检索、单样本或少样本学习、异常检测、相似性搜索、释义检测、聚类、分类等的基础。

嵌入可能难以扩展

然而，对于生产用例，嵌入可能难以扩展，这会导致昂贵的解决方案和高延迟。目前，许多最先进的模型产生具有 1024 维的嵌入，每个维度都以 float32 编码，即每个维度需要 4 个字节。因此，要对 2.5 亿个向量进行检索，您将需要大约 1TB 的内存！

下表概述了不同模型、维度大小、内存需求和成本。成本是根据 AWS 上 x2gd 实例每月每 GB 约 3.8 美元的估算费用计算的。

提高可扩展性

有几种方法可以应对扩展嵌入的挑战。最常见的方法是降维，例如 PCA。然而，经典的降维方法——如 PCA 方法——在与嵌入一起使用时往往表现不佳。

最近，OpenAI 使用的Matryoshka 表示学习（博客文章）(MRL) 也使得嵌入变得更经济。使用 MRL 时，只使用前 n 个嵌入维度。这种方法已经被一些开放模型采用，例如 nomic-ai/nomic-embed-text-v1.5 和 mixedbread-ai/mxbai-embed-2d-large-v1（对于 OpenAI 的 text-embedding-3-large，我们在 12 倍压缩下性能保持率为 93.1%。对于 nomic 的模型，我们在 3 倍压缩下保持 95.8% 的性能，在 6 倍压缩下保持 90% 的性能）。

然而，还有一种新的方法可以应对这一挑战；它不涉及降维，而是减少嵌入中每个单独值的大小：量化。我们关于量化的实验将表明，我们可以在保持大量性能的同时，显著加快计算速度并节省内存、存储和成本。让我们深入了解一下！

二值量化

与模型中降低权重精度的量化不同，嵌入的量化指的是对嵌入本身进行后处理。特别是，二值量化指的是将嵌入中的 float32 值转换为 1 位值，从而使内存和存储使用量减少 32 倍。

要将 float32 嵌入量化为二值，我们只需将归一化后的嵌入以 0 为阈值进行处理

$$f(x)= \begin{cases} 0 & \text{if } x\leq 0\\ 1 & \text{if } x \gt 0 \end{cases}$$

我们可以使用汉明距离来高效地检索这些二值嵌入。汉明距离是两个二值嵌入的比特位不同的位置数量。汉明距离越低，嵌入越接近；因此，文档越相关。汉明距离的一大优势是它可以用 2 个 CPU 周期轻松计算，从而实现极快的性能。

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

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

Sentence Transformers 中的二值量化

将维度为 1024 的嵌入量化为二值将产生 1024 位。在实践中，更常见的做法是将位存储为字节，因此当量化为二值嵌入时，我们使用 np.packbits 将位打包成字节。

因此，将维度为 1024 的 float32 嵌入量化后，会得到一个维度为 128 的 int8 或 uint8 嵌入。以下是使用 Sentence Transformers 生成量化嵌入的两种方法：

from sentence_transformers import SentenceTransformer

# 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",
)

或者

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")

# 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")

参考文献

• mixedbread-ai/mxbai-embed-large-v1
• SentenceTransformer.encode
• quantize_embeddings

在这里，您可以看到默认 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），如下图所示。这是通过使用大量的嵌入校准数据集来完成的。我们计算这些嵌入的范围，即每个嵌入维度的最小值和最大值。然后，我们计算步长（桶）来对每个值进行分类。

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

来源：https://qdrant.org.cn/articles/scalar-quantization/

通过对 int8 进行标量量化，我们降低了原始 float32 嵌入的精度，使每个值都用一个 8 位整数表示（小 4 倍）。请注意，这与二值量化的情况不同，在二值量化中，每个值由一个比特表示（小 32 倍）。

Sentence Transformers 中的标量量化

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

在实践中，建议为标量量化提供以下之一：

1. 大量要一次性量化的嵌入，或者
2. 每个嵌入维度的最小值和最大值范围，或者
3. 一个大的嵌入校准数据集，可以从中计算最小值和最大值范围。

如果以上情况均不满足，您将收到如下警告：基于 2 个嵌入计算 int8 量化桶。当使用从更多嵌入计算出的'ranges'或可用于计算桶的'calibration_embeddings'时，int8 量化更稳定。

以下是使用 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,
)

参考文献

• mixedbread-ai/mxbai-embed-large-v1
• SentenceTransformer.encode
• quantize_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. 使用量化后的查询搜索二值索引（4100 万个二值嵌入；5.2GB 内存/磁盘空间），以找到前 40 个文档。
4. 前 40 个文档从磁盘上的 int8 索引（4100 万个 int8 嵌入；0 字节内存，47.5GB 磁盘空间）中动态加载。
5. 使用 float32 查询和 int8 嵌入对前 40 个文档进行重排序，以获得前 10 个文档。
6. 前 10 个文档按分数排序并显示。

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

量化实验

我们在 MTEB 的检索子集上进行了实验，该子集包含 15 个基准测试。首先，我们使用 4 的 rescore_multiplier 检索了前 k（k=100）个搜索结果。因此，我们总共检索了 400 个结果，并对这前 400 个结果进行了重排序。对于 int8 性能，我们直接使用了点积而没有进行任何重排序。

从我们的量化实验结果中可以发现几个关键的趋势和好处。正如预期的那样，具有更高维度大小的嵌入模型通常会产生更高的每计算存储成本，但能实现最佳性能。然而，令人惊讶的是，量化到 int8 已经帮助 mxbai-embed-large-v1 和 Cohere-embed-english-v3.0 在存储使用量低于较小维度大小的基础模型的情况下实现了更高的性能。

当量化到二值模型时，量化的好处更加明显。在这种情况下，1024 维的模型仍然优于现在存储密集度高出 10 倍的基础模型，而 mxbai-embed-large-v1 在资源需求减少 32 倍后，甚至能够保持超过 96% 的性能。对于该模型，从 int8 到二值的进一步量化几乎没有带来额外的性能损失。

有趣的是，我们还可以看到 all-MiniLM-L6-v2 在二值量化上的性能优于 int8 量化。一个可能的解释是校准数据的选择。在 e5-base-v2 上，我们观察到 维度坍塌 的影响，这导致模型只使用潜在空间的一个子空间；在进行量化时，整个空间进一步坍塌，导致高性能损失。

这表明量化并非对所有嵌入模型都普遍适用。考虑现有的基准测试结果并进行实验以确定给定模型与量化的兼容性仍然至关重要。

重排序（Rescoring）的影响

在本节中，我们研究了重排序对检索性能的影响。我们基于 mxbai-embed-large-v1 评估结果。

二值重排序

使用二值嵌入，mxbai-embed-large-v1 在 MTEB 检索上保留了 92.53% 的性能。仅在不检索更多样本的情况下进行重排序，就将性能推至 96.45%。我们尝试将 rescore_multiplier 从 1 设置到 10，但没有观察到性能的进一步提升。这表明 top_k 搜索已经检索到了最佳候选者，而重排序只是适当地重新排列了这些好的候选者。

标量 (Int8) 重排序

我们还评估了带有 int8 重排序的 mxbai-embed-large-v1 模型，因为 Cohere 表明 Cohere-embed-english-v3.0 在 int8 量化下达到了 float32 模型 100% 的性能。对于这个实验，我们将 rescore_multiplier 设置为 [1, 4, 10]，并得到以下结果

从图中可以看出，更高的重排序乘数意味着量化后性能保留得更好。根据我们的结果推断，随着重排序乘数的持续增加，性能与 100% 的关系可能是双曲线的。重排序乘数为 4-5 已经能使用 int8 达到 99% 的卓越性能保留率。

检索速度

我们在 Google Cloud Platform a2-highgpu-4g 实例上测量了检索速度，使用了 1024 维的 mxbai-embed-large-v1 嵌入对整个 MTEB Retrieval 进行测试。对于 int8，我们使用了 USearch（版本 2.9.2），对于二值量化，我们使用了 Faiss（版本 1.8.0）。所有计算都在 CPU 上使用精确搜索完成。

如表所示，应用 int8 标量量化与全尺寸 float32 嵌入相比，平均速度提升了 3.66 倍。此外，二值量化平均实现了 24.76 倍的速度提升。对于标量和二值量化，即使在最坏的情况下也实现了非常显著的速度提升。

性能总结

实验结果、对资源使用、检索速度和检索性能的影响可以通过使用量化总结如下：

演示

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

亲自尝试

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

• 推荐的检索方法:
  • semantic_search_recommended.py：该脚本结合了二值搜索和标量重排序，与上述演示类似，实现了廉价、高效且高性能的检索。
• 用法:
  • semantic_search_faiss.py：该脚本通过使用 semantic_search_faiss 工具函数，展示了使用 FAISS 进行常规二值或标量量化、检索和重排序的用法。
  • semantic_search_usearch.py：该脚本通过使用 semantic_search_usearch 工具函数，展示了使用 USearch 进行常规二值或标量量化、检索和重排序的用法。
• 基准测试:
  • semantic_search_faiss_benchmark.py：该脚本使用 FAISS 对 float32 检索、二值检索+重排序、以及标量检索+重排序的检索速度进行了基准测试。它使用了 semantic_search_faiss 工具函数。我们的基准测试特别显示了 ubinary 的速度提升。
  • semantic_search_usearch_benchmark.py：该脚本使用 USearch 对 float32 检索、二值检索+重排序、以及标量检索+重排序的检索速度进行了基准测试。它使用了 semantic_search_usearch 工具函数。我们的实验显示，在较新的硬件上，特别是在 int8 方面，速度有很大提升。

未来工作

我们期待二值量化的进一步发展。举几个潜在的改进，我们怀疑可能存在比 int8 更小的标量量化空间，即使用 128 或 64 个桶而不是 256 个。

此外，我们很高兴地看到，嵌入量化与 Matryoshka 表示学习 (MRL) 完全正交。换句话说，可以将 MRL 嵌入从例如 1024 缩小到 128（通常对应于 2% 的性能下降），然后应用二值或标量量化。我们怀疑这可以使检索速度提高高达 32 倍，而质量下降约 3%，或者速度提高高达 256 倍，而质量下降约 10%。

最后，我们认识到使用嵌入量化的检索也可以与单独的重排序模型相结合。我们设想一个由二值搜索、标量 (int8) 重排序和交叉编码器重排序组成的 3 步流程，可以在低延迟、低内存使用、低磁盘空间和低成本的情况下实现最先进的检索性能。

致谢

这个项目得以实现，得益于我们与 mixedbread.ai 和 SentenceTransformers 库的合作，该库使您可以轻松创建句子嵌入并对其进行量化。如果您想在项目中使用量化嵌入，现在您知道该怎么做了！

引用

@article{shakir2024quantization,
  author       = { Aamir Shakir and
                   Tom Aarsen and
                   Sean Lee
                 },
  title = { Binary and Scalar Embedding Quantization for Significantly Faster & Cheaper Retrieval },
  journal = {Hugging Face Blog},
  year = {2024},
  note = {https://huggingface.co/blog/embedding-quantization},
}

资源

• mixedbread-ai/mxbai-embed-large-v1
• SentenceTransformer.encode
• quantize_embeddings
• Sentence Transformers 文档 - 嵌入量化
• https://txt.cohere.com/int8-binary-embeddings/
• https://qdrant.org.cn/documentation/guides/quantization
• https://zilliz.com/learn/scalar-quantization-and-product-quantization