用於顯著提高檢索速度和降低成本的二進位制和標量嵌入量化

HuggingFace發表於2024-04-12

我們引入了嵌入量化的概念，並展示了它們對檢索速度、記憶體使用、磁碟空間和成本的影響。我們將討論理論上和實踐中如何對嵌入進行量化，然後介紹一個 演示，展示了 4100 萬維基百科文字的真實檢索場景。

目錄

  • 為什麼使用嵌入？
    • 嵌入可能難以擴充套件
  • 提高可擴充套件性
    • 二進位制量化
      • Sentence Transformers 中的二進位制量化
      • 向量資料庫中的二進位制量化
    • 標量(int8)量化
      • Sentence Transformers 中的標量量化
      • 向量資料庫中的標量量化
    • 結合二進位制和標量量化
    • 量化實驗
    • 重打分的影響
      • 二進位制重打分
      • 標量（Int8）重打分
      • 檢索速度
    • 效能彙總
    • 演示
    • 自己嘗試
    • 未來工作
    • 致謝
    • 引用
    • 參考文獻

為什麼使用嵌入？

嵌入是自然語言處理中最多樣化的工具之一，支援各種設定和使用場景。本質上，嵌入是對更復雜物件 (如文字、影像、音訊等) 的數值表示。具體來說，這些物件被表示為 n 維向量。

在轉換了複雜物件之後，你可以透過計算相應嵌入的相似性來確定它們的相似性！這對於許多使用場景至關重要: 它為推薦系統、檢索、單次學習或少樣本學習、異常檢測、相似性搜尋、釋義檢測、聚類、分類等提供了基礎。

嵌入可能難以擴充套件

但是，當我們在實際應用中使用嵌入時，可能會遇到一些問題。比如，現在很多先進的模型產生的嵌入都是 1024 維的，每個維度需要 4 位元組的空間來儲存 (float 32 編碼)。如果你要處理 2.5 億個這樣的向量，那就需要大約 1TB 的記憶體，這既花錢又可能導致處理速度變慢。

下表展示了一些不同的模型，它們的維度大小、需要的記憶體量以及相應的成本。成本是按照 AWS 上一種叫做 x2gd 的例項來估算的，大概每個月每 GB 需要 3.8 美元。

嵌入維數 模型樣例 100M 嵌入 250M 嵌入 1B 嵌入
384 all-MiniLM-L6-v2
bge-small-en-v1.5		 143.05GB
$543 / mo | 357.62GB
$1,358 / mo		 1430.51GB
$5,435 / mo
768 all-mpnet-base-v2
bge-base-en-v1.5
jina-embeddings-v2-base-en
nomic-embed-text-v1		 286.10GB
$1,087 / mo|715.26GB
$2,717 / mo		 2861.02GB
$10,871 / mo
1024 bge-large-en-v1.5
mxbai-embed-large-v1
Cohere-embed-english-v3.0		 381.46GB
$1,449 / mo|953.67GB
$3,623 / mo		 3814.69GB
$14,495 / mo
1536 OpenAI text-embedding-3-small 572.20GB
$2,174 / mo|1430.51GB
$5,435 / mo		 5722.04GB
$21,743 / mo
3072 OpenAI text-embedding-3-large 1144.40GB
$4,348 / mo|2861.02GB
$10,871 / mo		 11444.09GB
$43,487 / mo

提高可擴充套件性

有幾種方法可以應對嵌入擴充套件的挑戰。最常見的方法是降維，比如使用 主成分分析 (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% 的效能。

然而，還有一種新的方法可以在這個挑戰上取得進展; 它不涉及降維，而是減少嵌入中每個個體值的尺寸大小: 量化。我們的量化實驗將展示，我們可以在顯著加快計算速度並節省記憶體、儲存和成本的同時，保持大量的效能。讓我們進一步瞭解一下吧！

二進位制量化

與在模型中減少權重精度的量化不同，嵌入的量化是指對嵌入本身進行的一個後處理步驟。特別是，二進位制量化指的是將嵌入中的 float32 值轉換為 1 bit ，從而在記憶體和儲存使用上實現 32 倍的減少。

要將 float32 嵌入量化為二進位制，我們只需將歸一化的嵌入在 0 處進行閾值處理:

\[ f(x)= \begin{cases} 0 & \text{如果 } x\leq 0\\ 1 & \text{如果 } x \gt 0 \end{cases} \]

我們可以使用漢明距離來高效地檢索這些二進位制嵌入。漢明距離是指兩個二進位制嵌入在位上不同的位置數量。漢明距離越低，嵌入越接近; 因此，文件的相關性越高。漢明距離的一個巨大優勢是它可以用 2 個 CPU 週期輕鬆計算，允許極快的效能。

Yamada 等人 (2021) 引入了一個重打分步驟，他們稱之為 rerank ，以提高效能。他們提議可以使用點積將 float32 查詢嵌入與二進位制文件嵌入進行比較。在實踐中，我們首先使用二進位制查詢嵌入和二進位制文件嵌入檢索 rescore_multiplier * top_k 的結果——即雙二進位制檢索的前 k 個結果的列表——然後使用 float32 查詢嵌入對這個二進位制文件嵌入列表進行重打分。

透過應用這種新穎的重打分步驟，我們能夠在減少記憶體和磁碟空間使用 32 倍的同時，保留高達 ~96% 的總檢索效能，並使檢索速度提高多達 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-v1SentenceTransformer.encodequantize_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 是
USearch 是
Vespa AI 是
Milvus 是
Qdrant 二進位制量化
Weaviate 二進位制量化

標量 (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 生成標量量化嵌入:

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-v1SentenceTransformer.encodequantize_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
USearch 是
Vespa AI 是
OpenSearch 是
ElasticSearch 是
Milvus IVF_SQ8
Qdrant 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 個文件進行排序並顯示。

透過這種方法，我們為索引使用了 5.2GB 記憶體和 52GB 磁碟空間。這比通常的檢索所需的 200GB 記憶體和 200GB 磁碟空間要少得多。尤其是當你進一步擴充套件時，這將顯著減少延遲和成本。

量化實驗

我們在 MTEB 的檢索子集上進行了實驗，該子集包含 15 個基準測試。首先，我們使用 rescore_multiplier 為 4 來檢索前 k (k=100) 個搜尋結果。因此，我們總共檢索了 400 個結果，並對這前 400 個結果進行了重打分。對於 int8 效能，我們直接使用了點積，而沒有進行任何重打分。

模型 嵌入維度 250M 嵌入 MTEB 檢索（NDCG@10） 預設效能的百分比
開源模型
mxbai-embed-large-v1: float32 1024 953.67GB
$3623 / mo		 54.39 100%
mxbai-embed-large-v1: int8 1024 238.41GB
$905 / mo		 52.79 97%
mxbai-embed-large-v1: binary 1024 29.80GB
$113.25 / mo		 52.46 96.45%
e5-base-v2: float32 768 286.10GB
$1087 / mo		 50.77 100%
e5-base-v2: int8 768 178.81GB
$679 / mo		 47.54 94.68%
e5-base-v2: binary 768 22.35GB
$85 / mo		 37.96 74.77%
all-MiniLM-L6-v2: float32 384 357.62GB
$1358 / mo		 41.66 100%
all-MiniLM-L6-v2: int8 384 89.40GB
$339 / mo		 37.82 90.79%
all-MiniLM-L6-v2: binary 384 11.18GB
$42 / mo		 39.07 93.79%
專有模型
Cohere-embed-english-v3.0: float32 1024 953.67GB
$3623 / mo		 55.0 100%
Cohere-embed-english-v3.0: int8 1024 238.41GB
$905 / mo		 55.0 100%
Cohere-embed-english-v3.0: binary 1024 29.80GB
$113.25 / mo		 52.3 94.6%

從我們的量化實驗結果中，可以識別出幾個關鍵趨勢和好處。正如預期的那樣，維度更高的嵌入模型通常每計算生成的儲存成本更高，但能實現最佳效能。然而，令人驚訝的是，量化到 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 上，我們觀察到了 維度坍縮 效應，導致模型只使用潛在空間的子空間; 當進行量化時，整個空間進一步坍縮，導致效能損失很大。

這表明量化並不適用於所有嵌入模型。考慮現有基準測試結果並開展實驗以確定給定模型與量化的相容性仍然至關重要。

重打分的影響

在本節中，我們探討了重打分對檢索效能的影響。我們基於 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% 的效能保留。

檢索速度

我們使用 Google Cloud Platform 的 a2-highgpu-4g 例項，在整個 MTEB 檢索中測量了 mxbai-embed-large-v1 嵌入的檢索速度，該嵌入的維度為 1024。對於 int8 ，我們使用了 USearch (版本 2.9.2) 和二進位制量化 Faiss (版本 1.8.0)。所有計算都在 CPU 上使用精確搜尋完成。

量化 最小 均值 最大
float32 1x (baseline) 1x (baseline) 1x (baseline)
int8 2.99x speedup 3.66x speedup 4.8x speedup
binary 15.05x speedup 24.76x speedup 45.8x speedup

如表中所示，應用 int8 標量量化相比全尺寸 float32 嵌入實現了平均速度提升 3.66 倍。此外，二進位制量化實現了平均速度提升 24.76 倍。對於標量和二進位制量化，即使在最壞的情況下，也實現了非常顯著的速度提升。

效能彙總

量化在資源使用、檢索速度和檢索效能方面的實驗結果和影響可以總結如下:

float32 int8/uint8 binary/ubinary
記憶體和索引空間節省 1x 精確 4x 精確 32x
檢索速度 1x 多達 4x 多達 45x
預設效能百分比 100% ~99.3% ~96%

演示

以下 演示 展示了透過結合二進位制搜尋和標量 ( int8 ) 重打分來提高檢索效率。該解決方案需要 5GB 的記憶體用於二進位制索引和 50GB 的磁碟空間用於二進位制和標量索引，這比常規的 float32 檢索所需的 200GB 記憶體和磁碟空間要少得多。此外，檢索速度也更快。

自己嘗試

以下指令碼可用於實驗性地進行檢索和其他用途的嵌入量化。它們分為三個類別:

  • 推薦檢索:

    • semantic_search_recommended.py: 此指令碼結合了二進位制搜尋和標量重打分，與上面的演示類似，以實現廉價、高效且效能良好的檢索。

  • 使用:

    • semantic_search_faiss.py: 此指令碼展示了使用 FAISS 的常規二進位制或標量量化、檢索和重打分的使用方式，透過使用 semantic_search_faiss 實用函式。
    • semantic_search_usearch.py: 此指令碼展示了使用 USearch 的常規二進位制或標量量化、檢索和重打分的使用方式，透過使用 semantic_search_usearch 實用函式。

  • 基準測試:

    • semantic_search_faiss_benchmark.py: 此指令碼包括了對 float32 檢索、二進位制檢索加重打分和標量檢索加重打分的檢索速度基準測試，使用 FAISS。它使用了 semantic_search_faiss 實用函式。我們的基準測試特別顯示了 ubinary 的速度提升。
    • semantic_search_usearch_benchmark.py: 此指令碼包括了對 float32 檢索、二進位制檢索加重打分和標量檢索加重打分的檢索速度基準測試，使用 USearch。它使用了 semantic_search_usearch 實用函式。我們的實驗在新硬體上顯示了巨大的速度提升，特別是對於 int8 。

未來工作

我們期待二進位制量化技術的進一步發展。要提到的一些潛在改進，我們懷疑可能還有比 int8 更小的標量量化空間，即使用 128 或 64 個桶而不是 256 個。

此外，我們也很興奮地發現，嵌入量化與 Matryoshka 表徵學習 (MRL) 完全垂直。換句話說，可以將 MRL 嵌入從例如 1024 減少到 128 (通常與 2% 的效能降低相對應)，然後應用二進位制或標量量化。我們懷疑這可能會將檢索速度提高多達 32 倍，質量降低約 3%，或者質量降低約 10% 時檢索速度提高多達 256 倍。

最後，我們意識到，使用嵌入量化進行檢索可以與一個獨立的重新排序模型結合起來使用。我們設想了一個三步流水線: 首先進行二進位制搜尋，然後對結果進行標量 (int8) 重打分，最後使用交叉編碼模型進行重新排序。這樣的流水線可以實現高效的檢索效能，同時降低延遲、記憶體使用、磁碟空間和成本。

致謝

這個專案得益於我們與 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-v1SentenceTransformer.encodequantize_embeddings

  • Sentence Transformers docs - Embedding Quantization
  • https://txt.cohere.com/int8-binary-embeddings/
  • https://qdrant.tech/documentation/guides/quantization
  • https://zilliz.com/learn/scalar-quantization-and-product-quantization

英文原文: https://hf.co/blog/embedding-quantization
原文作者: Aamir Shakir, Tom Aarsen, SeanLee
譯者: innovation64

相關文章