Sentence Transformersを使用して400倍高速なStatic Embedding Modelをトレーニングする

TL;DR

Hugging Faceは、標準的なTransformerベースの埋め込みの品質を少なくとも85%維持しつつ、CPU上で100倍〜400倍高速に動作する静的埋め込みモデル（static embedding models）のトレーニングレシピを公開しました。また、2つの具体的なモデル（英語検索用の static-retrieval-mrl-en-v1 と多言語類似度用の static-similarity-mrl-multilingual-v1）を、トレーニングスクリプト、評価結果、Weights & Biasesのログと共に公開しました。

手法

このアプローチは、対照学習（contrastive learning）を MultipleNegativesRankingLoss およびオプションのMatryoshka Representation Learningと組み合わせることで、アテンションベースのエンコーディングの代わりに単純なトークン埋め込みのルックアップを行う静的エンコーダをトレーニングします。これにより、品質の低下を最小限に抑えつつ、数桁の高速化を実現します。

トレーニングの詳細

要件

トレーニングには、dataset、loss function、training arguments、evaluator、trainer のコンポーネントを備えた Sentence Transformers ライブラリを使用します。

モデルのインスピレーション

ターゲットとしたのは2つのモデルです。1つは英語専用の検索モデル、もう1つは多言語の汎用類似度モデルであり、どちらも BERTベースのトークナイザー（bert-base-uncased または bert-base-multilingual-uncased）をラップした StaticEmbedding モジュールと、1024の埋め込み次元を使用して初期化されています。

データセットの選択

英語検索には、gooaq、msmarco (triplet)、squad、s2orc (title-abstract-pair)、allnli (triplet)、paq、trivia_qa、msmarco_10m、swim_ir (en)、pubmedqa (triplet-20)、miracl (en-triplet-all)、mldr (en-triplet-all)、mr_tydi (en-triplet-all) を含む30のデータセットが選択されました。多言語類似度については、並列文を持つデータセット（wikititles、tatoeba、talks、europarl、global_voices、jw300、muse、wikimatrix、opensubtitles）と、ポジティブペアのデータセット（stackexchange-duplicates、quora-duplicates、wikianswers-duplicates、allnli、simple_wiki、altlex、flickr30k_captions、coco_captions、nli_for_simcse、negation）を組み合わせ、計30のトレーニングデータセットを使用しました。

損失関数

MultipleNegativesRankingLoss が選ばれた理由は、バッチサイズ2048が RTX 3090 に収まり、CachedMultipleNegativesRankingLoss のオーバーヘッドや GISTEmbedLoss におけるガイドモデルによる低速化を回避できるためです。その上に、次元 [32, 64, 128, 256, 512, 1024] で MatryoshkaLoss が適用されました。

トレーニング引数

両モデルで使用された設定：num_train_epochs=1、per_device_train_batch_size=2048、per_device_eval_batch_size=2048、learning_rate=2e-1、warmup_ratio=0.1、bf16=True、batch_sampler=BatchSamplers.NO_DUPLICATES、multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL、eval_strategy=steps、eval_steps=250 (英語) または 1000 (多言語)、save_strategy=steps、save_steps は eval_steps に一致、save_total_limit=2、logging_steps は eval_steps に一致、logging_first_step=True、および実行ごとの出力ディレクトリ。

評価器

英語検索モデルには、ゼロショット検索評価のために NanoBEIREvaluator を使用しました。多言語モデルは、評価のために MTEB タスク（STS、Classification、Pair Classification）に依存しています。

ハードウェア

トレーニングは、RTX 3090 GPU、i7-13700K CPU、および 32 GB RAM で実行されました。

全体的なトレーニングスクリプト

提供されたスクリプトは、データセットをロードし、StaticEmbedding モデルをインスタンス化し、MatryoshkaLoss で損失を構成し、トレーニング引数を設定し、オプションで評価器を実行し、SentenceTransformerTrainer でトレーニングを行い、最終的なモデルを保存します。英語検索スクリプトは17.8時間かかり、2.6 kWhを消費し、1 kgのCO₂を排出しました。多言語スクリプトは3.1時間かかり、0.5 kWhを消費し、0.2 kgのCO₂を排出しました。

使用方法

両モデルは、モデル名と device="cpu" を指定して SentenceTransformer を介してロードされます。推論は標準的な Sentence Transformers と同一です：model.encode は埋め込みを返し、model.similarity はコサイン類似度を計算します。truncate_dim 引数により、Matryoshka スタイルの次元削減が可能です（例：truncate_dim=256）。これらのモデルは、LangChain、LlamaIndex、Haystack、txtai でそのまま動作します。

パフォーマンス

英語検索

NanoBEIR において、static-retrieval-mrl-en-v1 は NDCG@10 で 0.5032 を達成しました。これは all-mpnet-base-v2 (0.5757) のスコアの 87.4% です。CPU では毎秒 107,419.51 文を処理し、all-mpnet-base-v2 (270.40 文/秒) よりも 397倍高速です。GPU では毎秒 97,171.47 文を処理し、all-mpnet-base-v2 (4043.13 文/秒) よりも 24倍高速です。Matryoshka 評価では、次元を 512 に半分に減らすと、NDCG@10 はわずか 1.47% 低下するだけです (0.5032 → 0.4957)。

多言語類似度

multilingual-e5-small と比較して、static-similarity-mrl-multilingual-v1 は STS で 92.3%、Pair Classification で 95.52%、Classification で 86.52% のスコアを記録しました。これは multilingual-e5-small よりも CPU で約 125倍、GPU で 10倍高速です。Matryoshka 評価では、次元を 256 (4分の1) に削減しても、英語 STS のパフォーマンスの低下はわずか 0.56% であることが示されています。

結論

提示されたレシピでトレーニングされた静的埋め込みモデルは、一般的な Transformer ベースの埋め込みの品質を少なくとも 85% 維持しながら、CPU で 100倍〜400倍、GPU で 10倍〜25倍の高速化を実現します。公開されたモデルにより、精度低下を最小限に抑えた効率的なオンデバイス、ブラウザ内、およびエッジコンピューティングのユースケースが可能になります。

次のステップ

ユーザーは、既存の Sentence Transformer モデルを static-retrieval-mrl-en-v1 または static-similarity-mrl-multilingual-v1 に置き換えるか、タスク固有のデータで独自の静的埋め込みをトレーニングできます。潜在的な改善策には、ハードネガティブマイニング、モデルスーピング、カリキュラム学習、ガイド付きインバッチ・ネガティブ・フィルタリング、シード最適化ランダム初期化、トークナイザーの再トレーニング、CachedMultipleNegativesRankingLoss による勾配キャッシュ、および大規模なエンコーダからのモデル蒸留が含まれます。