二值和標量嵌入量化：實現更快、更經濟的檢索

我們引入嵌入量化的概念，並展示其對檢索速度、記憶體使用、磁碟空間和成本的影響。我們將討論嵌入在理論上和實踐中如何進行量化，然後介紹一個演示，該演示展示了對 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.tech/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.tech/documentation/guides/quantization
• https://zilliz.com/learn/scalar-quantization-and-product-quantization