埋め込み量子化:高速・低コスト検索のためのバイナリおよびスカラー手法
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 段階のパイプラインにより、両者の長所を最大限に活かすことができます:
- 高品質モデル(例:
mxbai-embed-large-v1)でクエリをエンコードする。 - クエリをバイナリに量子化し、バイナリインデックスで検索する(41 M の Wikipedia パッセージで約 5 GB)。
- ディスク上に保存された int8 インデックスから上位 k 件の候補をロードする(約 48 GB)。
- 元の
float32クエリを使用して、int8 埋め込みに対してこれらの候補を再スコアリングする。 - 最終的な上位 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 | 1× | 1× | 1× |
| int8 | 2.99× | 3.66× | 4.8× |
| binary | 15.05× | 24.76× | 45.8× |
| したがって、バイナリ量子化は桁違いのレイテンシ削減を提供します。 |
トレードオフの概要
| 指標 | float32 | int8/uint8 | binary/ubinary |
|---|---|---|---|
| メモリ & インデックスサイズ | 1× | 4× smaller | 32× smaller |
| 検索速度 | 1× | 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}
}