用于显著提高检索速度和降低成本的二进制和标量嵌入量化

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

• 演示地址 https://hf.co/spaces/sentence-transformers/quantized-retrieval

为什么使用嵌入？

嵌入是自然语言处理中最多样化的工具之一，支持各种设置和使用场景。本质上，嵌入是对更复杂对象 (如文本、图像、音频等) 的数值表示。具体来说，这些对象被表示为 n 维向量。

在转换了复杂对象之后，你可以通过计算相应嵌入的相似性来确定它们的相似性！这对于许多使用场景至关重要: 它为推荐系统、检索、单次学习或少样本学习、异常检测、相似性搜索、释义检测、聚类、分类等提供了基础。

嵌入可能难以扩展

但是，当我们在实际应用中使用嵌入时，可能会遇到一些问题。比如，现在很多先进的模型产生的嵌入都是 1024 维的，每个维度需要 4 字节的空间来存储 (float 32 编码)。如果你要处理 2.5 亿个这样的向量，那就需要大约 1TB 的内存，这既花钱又可能导致处理速度变慢。

下表展示了一些不同的模型，它们的维度大小、需要的内存量以及相应的成本。成本是按照 AWS 上一种叫做 x2gd 的实例来估算的，大概每个月每 GB 需要 3.8 美元。

• all-MiniLM-L6-v2 https://hf.co/sentence-transformers/all-MiniLM-L6-v2
• bge-small-en-v1.5 https://hf.co/BAAI/bge-small-en-v1.5
• all-mpnet-base-v2 https://hf.co/sentence-transformers/all-mpnet-base-v2
• bge-base-en-v1.5 https://hf.co/BAAI/bge-base-en-v1.5
• jina-embeddings-v2-base-en https://hf.co/jinaai/jina-embeddings-v2-base-en
• nomic-embed-text-v1 https://hf.co/nomic-ai/nomic-embed-text-v1
• bge-large-en-v1.5 https://huggingface.co/BAAI/bge-large-en-v1.5
• mxbai-embed-large-v1 https://hf.co/mixedbread-ai/mxbai-embed-large-v1
• Cohere-embed-english-v3.0 https://txt.cohere.com/introducing-embed-v3/
• OpenAI text-embedding-3-small https://openai.com/blog/new-embedding-models-and-api-updates
• OpenAI text-embedding-3-large https://openai.com/blog/new-embedding-models-and-api-updates

提高可扩展性

有几种方法可以应对嵌入扩展的挑战。最常见的方法是降维，比如使用 主成分分析 (PCA)。然而，传统的降维方法——比如 PCA ——在处理嵌入时往往效果不佳。最近，有关于  Matryoshka 表征学习 (MRL) 的新闻 (博客)，这种方法由 OpenAI 使用，允许更经济的嵌入。使用 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% 的性能。

• 主成分分析 (PCA) https://en.wikipedia.org/wiki/Principal_component_analysis
• Matryoshka 表征学习 https://arxiv.org/abs/2205.13147
• 博客 https://huggingface.co/blog/matryoshka
• OpenAI https://openai.com/blog/new-embedding-models-and-api-updates
• nomic-ai/nomic-embed-text-v1.5 https://hf.co/nomic-ai/nomic-embed-text-v1.5
• mixedbread-ai/mxbai-embed-2d-large-v1 https://hf.co/mixedbread-ai/mxbai-embed-2d-large-v1

然而，还有一种新的方法可以在这个挑战上取得进展; 它不涉及降维，而是减少嵌入中每个个体值的尺寸大小: 量化 。我们的量化实验将展示，我们可以在显著加快计算速度并节省内存、存储和成本的同时，保持大量的性能。让我们进一步了解一下吧！

二进制量化

与在模型中减少权重精度的量化不同，嵌入的量化是指对嵌入本身进行的一个后处理步骤。特别是，二进制量化指的是将嵌入中的 float32 值转换为 1 bit ，从而在内存和存储使用上实现 32 倍的减少。

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

我们可以使用汉明距离来高效地检索这些二进制嵌入。汉明距离是指两个二进制嵌入在位上不同的位置数量。汉明距离越低，嵌入越接近; 因此，文档的相关性越高。汉明距离的一个巨大优势是它可以用 2 个 CPU 周期轻松计算，允许极快的性能。

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

• Yamada 等人 (2021) https://arxiv.org/abs/2106.00882

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

Sentence Transformers 中的二进制量化

将一个维度为 1024 的嵌入量化为二进制将得到 1024 比特。实际上，将比特存储为字节要常见得多，因此当我们量化为二进制嵌入时，我们使用 np.packbits 将比特打包成字节。

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

• Sentence Transformers https://sbert.net/

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 数据类型方面的差异:

>>> 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 数据格式将嵌入量化为二进制。这可能取决于你的向量库/数据库的要求。

向量数据库中的二进制量化

• Faiss https://github.com/facebookresearch/faiss/wiki/Binary-indexes
• USearch https://github.com/unum-cloud/usearch
• Vespa AI https://docs.vespa.ai/en/reference/schema-reference.html
• Milvus https://milvus.io/docs/index.md
• Qdrant 二进制量化 https://qdrant.tech/documentation/guides/quantization/#binary-quantization
• Weaviate 二进制量化 https://weaviate.io/developers/weaviate/configuration/bq-compression

标量 (int8) 量化

我们使用标量量化过程将 float32 嵌入转换为 int8 。这涉及到将 float32 值的连续范围映射到可以表示 256 个不同级别 (从 -128 到 127) 的 int8 值的离散集合，如下面的图像所示。这是通过使用大量的嵌入校准数据集来完成的。我们计算这些嵌入的范围，即每个嵌入维度的 min 和 max 。从这里，我们计算将每个值分类的步骤 (桶)。

为了进一步提高检索性能，你可以可选地应用与二进制嵌入相同的重打分步骤。重要的是要注意，校准数据集极大地影响性能，因为它定义了量化桶。

Source: https://qdrant.tech/articles/scalar-quantization/

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

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. 大意是如果你只使用很少量的嵌入 (在这个例子中是 2 个嵌入) 来计算这些量化桶，那么量化可能不会那么稳定或准确，因为少量的数据可能无法很好地代表整个数据分布。因此，如果你有一个很大的数据集来计算这些范围，或者有一个校准数据集，那么你可以得到更好的量化结果。

请看下面如何使用 Sentence Transformers 生成标量量化嵌入:

• Sentence Transformers https://sbert.net/

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 数据类型方面的差异:

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

向量数据库中的标量量化

• Faiss: IndexHNSWSQ https://faiss.ai/cpp_api/struct/structfaiss_1_1IndexHNSWSQ.html
• USearch https://github.com/unum-cloud/usearch
• Vespa AI https://docs.vespa.ai/en/reference/tensor.html
• OpenSearch https://opensearch.org/docs/latest/field-types/supported-field-types/knn-vector
• ElasticSearch https://www.elastic.co/de/blog/save-space-with-byte-sized-vectors
• Milvus: IVF_SQ8 https://milvus.io/docs/index.md
• Qdrant: Scalar Quantization https://qdrant.tech/documentation/guides/quantization/#scalar-quantization

结合二进制和标量量化

结合二进制和标量量化可以兼得两者的优点: 二进制嵌入的极快速度和标量嵌入在重打分后的优良性能的保留。请查看下面的演示，这是一个涉及维基百科 4100 万文本的真实实现。该设置的流程如下:

1. 使用 mixedbread-ai/mxbai-embed-large-v1 SentenceTransformer 模型对查询进行嵌入。
2. 使用 sentence-transformers 库中的 quantize_embeddings 函数将查询量化为二进制。
3. 使用量化查询在二进制索引 (4100 万个二进制嵌入; 5.2GB 内存/磁盘空间) 中搜索前 40 个文档。
4. 从磁盘上的 int8 索引 (4100 万个 int8 嵌入; 0 字节内存，47.5GB 磁盘空间) 动态加载前 40 个文档。
5. 使用 float32 查询和 int8 嵌入对前 40 个文档进行重打分，以获得前 10 个文档。
6. 按分数对前 10 个文档进行排序并显示。

• mixedbread-ai/mxbai-embed-large-v1 https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1

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

量化实验

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

• MTEB https://huggingface.co/spaces/mteb/leaderboard

• mxbai-embed-large-v1 https://hf.co/mixedbread-ai/mxbai-embed-large-v1
• e5-base-v2 https://hf.co/intfloat/e5-base-v2
• all-MiniLM-L6-v2 https://hf.co/sentence-transformers/all-MiniLM-L6-v2
• Cohere-embed-english-v3.0 https://txt.cohere.com/introducing-embed-v3/

从我们的量化实验结果中，可以识别出几个关键趋势和好处。正如预期的那样，维度更高的嵌入模型通常每计算生成的存储成本更高，但能实现最佳性能。然而，令人惊讶的是，量化到 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 上，我们观察到了 维度坍缩 效应，导致模型只使用潜在空间的子空间; 当进行量化时，整个空间进一步坍缩，导致性能损失很大。

• 维度坍缩 https://arxiv.org/abs/2110.09348

这表明量化并不适用于所有嵌入模型。考虑现有基准测试结果并开展实验以确定给定模型与量化的兼容性仍然至关重要。

重打分的影响

在本节中，我们探讨了重打分对检索性能的影响。我们基于 mxbai-embed-large-v1 评估了结果。

二进制重打分

使用二进制嵌入，mxbai-embed-large-v1 在 MTEB 检索上保留了 92.53% 的性能。仅进行重打分而无需检索更多样本，性能提升到了 96.45%。我们实验设置了 rescore_multiplier 从 1 到 10，但没有观察到进一步的性能提升。这表明 top_k 搜索已经检索到了最顶级的候选项，而重打分则正确地重新排列了这些好的候选项。

标量 (Int8) 重打分

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

从图表中我们可以看到，更高的重打分乘数意味着量化后性能的更好保留。从我们的结果推断，我们假设这种关系可能是双曲线的，随着重打分乘数的增加，性能接近 100%。使用 int8 时，重打分乘数为 4-5 已经导致令人瞩目的 99% 的性能保留。