埋め込み量子化:高速・低コスト検索のためのバイナリおよびスカラー手法

TL;DR: Hugging Face はバイナリおよびスカラー(int8)埋め込み量子化を導入し、埋め込みをそれぞれ 32 倍または 4 倍に圧縮し、メモリとストレージコストを大幅に削減し、検索速度を最大 45 倍に向上させ、元の性能の 96%〜99% を維持します。

埋め込みが重要な理由とスケーリング方法

埋め込みはテキスト、画像、音声、その他のデータを高次元ベクトルに変換し、類似検索、レコメンデーション、クラスタリング、そして多くの下流 NLP タスクを可能にします。最先端モデルはしばしば 1024 次元の float32 ベクトルを出力し、次元あたり 4 バイトが必要です。250 M 個のベクトルを保存すると約 1 TB の RAM を消費し、月額数千ドル規模のクラウド費用が発生します。ブログでは複数の人気モデルのコストを定量化し、1024 次元モデルが AWS x2gd インスタンスで月額 $3,600 超になることを示しています。

量子化と次元削減の比較

従来のスケーリング手法は次元削減(例:PCA)や Matryoshka Representation Learning(MRL)を使用し、次元を削減しますが、性能が低下することがあります。埋め込み量子化は、モデルが埋め込みを生成した後に各次元の 精度 を低減させ、低コスト検索への補完的なアプローチを提供します。

バイナリ量子化

バイナリ量子化は各 float32 値をゼロで閾値処理し、1 ビットに変換します。これによりストレージが 32 倍削減され(例:1024 次元ベクトルは 1024 ビットとなり、128 バイトにパックされます)。検索はハミング距離を使用し、わずか 2 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 データ型で、元の float32 埋め込みが 8 192 バイトであるのに対し、256 バイトしか占有しません。

ベクトルデータベースでのサポート

バイナリインデックスは 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 で(直接または間接的に)サポートされています。

バイナリとスカラー量子化の組み合わせ

2 段階のパイプラインにより、両者の長所を最大限に活かすことができます:

  1. 高品質モデル(例:mxbai-embed-large-v1)でクエリをエンコードする。
  2. クエリをバイナリに量子化し、バイナリインデックスで検索する(41 M の Wikipedia パッセージで約 5 GB)。
  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 量子化は、ストレージを 4 倍削減しながら、性能の >94 % を維持することが多いです。
  • バイナリ量子化は大次元モデルで約 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× smaller 32× smaller
検索速度 up to 4× faster up to 45× faster
性能保持率 100 % ~99 % ~96 %

デモと実用スクリプト

ライブデモ(link)は、5 GB の RAM と 52 GB のディスクで 41 M の Wikipedia パッセージの検索を示し、上記の速度向上を実現しています。ブログでは、すぐに実行できるスクリプトを 3 つのカテゴリで提供しています:

  • 推奨検索 – バイナリ検索と int8 再スコアリングを組み合わせます。
  • 使用方法 – 量子化された埋め込みで semantic_search_faiss または semantic_search_usearch を呼び出す方法を示します。
  • ベンチマーク – 各量子化モードの速度と精度を測定します。

今後の方向性

  • サブ int8 量子化(例:4 ビットや 2 ビットバケット)を検討し、さらに高い圧縮率を目指す。
  • 量子化と Matryoshka Representation Learning を組み合わせ、まず次元を削減し、その後量子化することで、品質低下を抑えつつ 32×〜256× の速度向上を実現できる可能性があります。
  • バイナリ + int8 ステージの後に第 3 ステージのクロスエンコーダ再ランカーを統合し、低レイテンシ・低コストで最先端の検索を実現します。

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