嵌入量化:二元與標量技術實現更快、更便宜的檢索

TL;DR: Hugging Face 介紹了二元與標量(int8)嵌入量化,分別將嵌入壓縮 32 倍或 4 倍,顯著降低記憶體與儲存成本,並將檢索速度提升至最高 45 倍,同時保留 96%–99% 的原始效能。

為何嵌入重要以及它們的擴展方式

嵌入將文字、圖像、音訊及其他資料轉換為高維向量,使相似度搜尋、推薦、分群以及許多下游 NLP 任務成為可能。最先進的模型通常產生 1024 維的 float32 向量,每個維度需要 4 位元組。儲存 2.5 億個此類向量約佔用 1 TB 記憶體,導致每月數千美元的雲端費用。部落格針對多個熱門模型量化了這些成本,顯示 1024 維模型在 AWS x2gd 實例上每月可能超過 $3,600。

量化 vs. 維度縮減

傳統的擴展方法使用維度縮減(例如 PCA)或 Matryoshka Representation Learning(MRL),會截斷維度但可能影響效能。嵌入量化則是在模型產生嵌入後降低每個維度的精度,提供一條互補的低成本檢索路徑。

二元量化

二元量化透過以零為閾值將每個 float32 值轉換為單一位元。這可將儲存空間縮減 32 倍(例如 1024 維向量變為 1024 位元,壓縮至 128 位元組)。檢索使用 Hamming 距離,只需兩個 CPU 週期即可計算,帶來巨大的加速。

在 Sentence‑Transformers 中的實作

from sentence_transformers import SentenceTransformer
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
# Direct binary encoding
binary_embeddings = model.encode(
    ["I am driving to the lake.", "It is a beautiful day."],
    precision="binary",
)

產生的 binary_embeddings 形狀為 (2, 128),資料型別為 int8,佔用 256 位元組,而原始的 float32 嵌入則佔 8 192 位元組。

向量資料庫的支援

二元索引已在 Faiss、USearch、Vespa AI、Milvus、Qdrant 與 Weaviate 中提供,可直接取代現有的管線。

標量(int8)量化

標量量化將每個維度的連續 float32 範圍映射為 256 個離散的 int8 級別(‑128 到 127)。這可將儲存空間縮減 4 倍,同時保留比二元更細緻的粒度。需要在大型嵌入集合上進行校準,以計算每個維度的最小/最大範圍。

在 Sentence‑Transformers 中的實作

from sentence_transformers import SentenceTransformer, quantize_embeddings
from datasets import load_dataset
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
corpus = load_dataset("nq_open", split="train[:1000]")['question']
calibration_embeddings = model.encode(corpus)
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,
)

int8_embeddings 保持原始 1024 維的形狀,但僅使用 2 048 位元組。

向量資料庫的支援

標量量化在 Faiss(IndexHNSWSQ)、USearch、Vespa AI、OpenSearch、ElasticSearch、Milvus(IVF_SQ8)與 Qdrant 中皆有支援(直接或間接)。

結合二元與標量量化

一個兩階段的管線可以同時取得兩者的優勢:

  1. 使用高品質模型(例如 mxbai-embed-large-v1)對查詢進行編碼。
  2. 將查詢量化為二元並在二元索引中搜尋(約 5 GB,針對 4100 萬篇 Wikipedia 文章)。
  3. 從磁碟上儲存的 int8 索引載入前 k 名候選(約 48 GB)。
  4. 使用原始 float32 查詢對 int8 嵌入重新打分。
  5. 返回最終的前 k 名結果。 此方法將記憶體需求降至約 5 GB,磁碟需求降至約 52 GB,遠低於完整精度檢索所需的約 200 GB。

實驗結果

檢索效能

模型 維度 儲存空間 (250 M) MTEB 檢索 NDCG@10 %(相對於 Float32)
mxbai-embed-large-v1 (float32) 1024 953.67 GB $3 623/mo 54.39 100 %
mxbai-embed-large-v1 (int8) 1024 238.41 GB $905/mo 52.79 97 %
mxbai-embed-large-v1 (binary) 1024 29.80 GB $113/mo 52.46 96.45 %
all-MiniLM-L6-v2 (binary) 384 11.18 GB $42/mo 39.07 93.79 %

主要觀察:

  • Int8 量化通常保留 >94 % 的效能,同時將儲存空間縮減 4 倍。
  • 二元量化對於大維度模型保留約 96 % 的效能,且在某些小模型(例如 all-MiniLM-L6-v2)上甚至優於 int8。
  • 效能因模型而異;校準資料品質與維度崩潰可能影響結果。

重新打分的影響

  • 二元重新打分(使用原始 float 查詢重新排序前 k 名二元結果)將效能從基線的 92.5 % 提升至 96.5 %。
  • 對於 int8,提升 rescore_multiplier(在重新打分前檢索更多候選)可提升保留率,在乘數為 4–5 時達到約 99 %。

檢索速度

在 GCP a2-highgpu-4g 僅 CPU 的精確搜尋上:

量化方式 最小加速倍數 平均加速倍數 最大加速倍數
float32
int8 2.99× 3.66× 4.8×
binary 15.05× 24.76× 45.8×
因此二元量化可提供量級的延遲降低。

折衷的總結

指標 float32 int8/uint8 binary/ubinary
記憶體與索引大小 4× 較小 32× 較小
檢索速度 最快 4× 加速 最快 45× 加速
效能保留 100 % ~99 % ~96 %

示範與實用腳本

一個線上示範(link)展示了在 41 M Wikipedia 文章上使用 5 GB 記憶體與 52 GB 磁碟的檢索,達成上述加速。部落格亦提供三類即用腳本:

  • 建議的檢索 – 結合二元搜尋與 int8 重新打分。
  • 使用方式 – 示範如何使用量化嵌入呼叫 semantic_search_faisssemantic_search_usearch
  • 基準測試 – 測量每種量化模式的速度與準確度。

未來方向

  • 探索次 int8 量化(例如 4 位元或 2 位元桶)以實現更緊密的壓縮。
  • 將量化與 Matryoshka Representation Learning 結合,先截斷維度再量化,可能在品質略有下降的情況下達到 32×–256× 的加速。
  • 在二元 + int8 階段之後加入第三階段的 cross‑encoder 重新排序器,以在低延遲與低成本下實現最先進的檢索。

Citation

@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}
}

Sources