Matryoshka 埋め込みモデルの紹介
Matryoshka 埋め込みモデルは、性能に大きな損失を伴わずに小さな次元へ切り詰め可能な埋め込みの作成を可能にします。このアプローチにより、実務者はストレージコスト、処理速度、精度要件に応じて埋め込みソリューションを動的にスケールさせることができます。
Matryoshka 埋め込みの仕組み
Matryoshka 埋め込みは Matryoshka Representation Learning (MRL) に基づいており、モデルが埋め込みベクトルの最初の次元に最も重要な情報を格納するよう促す手法です。ロシアの入れ子人形のように、モデルは埋め込みの小さく切り詰めたバージョンでも、下流タスクに有用な情報を十分に保持できるように訓練されます。
応用と利点
可変サイズの埋め込みは、AI 実務者に対して主に 2 つの利点を提供します。
- ショートリスト作成と再ランク付け: ユーザーは小さく切り詰めた埋め込みを使用して、初期の非常に効率的な「ショートリスト」検索を実行できます。候補セットが特定されたら、全次元の埋め込みを使用して結果を再ランク付けし、精度を向上させることができます。
- リソース最適化: Matryoshka モデルは、デプロイ環境の具体的な制約に応じて、ストレージコスト、処理速度、性能の間でカスタマイズ可能なトレードオフを可能にします。
Matryoshka 埋め込みモデルのトレーニング
理論的枠組み
標準的な埋め込みトレーニングでは、損失関数はフルサイズの埋め込みに適用されます。これに対し、Matryoshka のトレーニングでは、フルサイズの埋め込みだけでなく、さまざまな切り詰め次元（例: 768、512、256、128、64）における埋め込みに対しても損失値を計算します。これらの損失を合計して最終的な損失値を作成し、オプティマイザがモデルの重みを調整する際に使用します。これにより、最も重要な情報が「前方」にロードされます。
Sentence Transformers における実装
Sentence Transformers は MatryoshkaLoss を通じて Matryoshka モデルのサポートを実装しました。このラッパー損失関数は、ベース損失（例: CoSENTLoss）を埋め込みの複数の切り詰め部分に適用します。MatryoshkaLoss を使用したトレーニングは、トレーニング時間に顕著なオーバーヘッドをもたらしません。
Matryoshka 埋め込みモデルの使用
一般的な推論
推論は標準的な埋め込みモデルと同様に機能します。唯一の違いは、生成された埋め込みをオプションで小さな次元に切り詰められることです。埋め込みが元々正規化されていた場合、切り詰め後に再正規化して一貫性を保つ必要があります。
Sentence Transformers における実装
SentenceTransformer クラスを使用すると、ユーザーはモデル初期化時に truncate_dim 引数を指定できます。encode 関数は、出力された埋め込みを指定されたサイズに自動的に切り詰めます。
nomic-ai/nomic-embed-text-v1.5 のような特定のモデルでは、埋め込みを切り詰める前に F.layer_norm を適用するカスタムアーキテクチャが必要です。
パフォーマンス結果
STSBenchmark テストセットで Matryoshka モデル（tomaarsen/mpnet-base-nli-matryoshka）と標準モデル（tomaarsen/mpnet-base-nli）を比較した実験結果は、以下のことを示しています。
- 精度向上: Matryoshka モデルは、テストされたすべての次元において、より高い Spearman 類似度を達成しました。
- 性能保持: Matryoshka モデルの性能低下は標準モデルに比べてはるかに緩やかです。元のサイズの 8.3% に切り詰めても、Matryoshka モデルは最大性能の 98.37% を維持し、標準モデルは 96.46% でした。
これらの結果は、Matryoshka 埋め込みが性能に顕著な影響を与えることなく、ストレージを大幅に削減し、検索速度を向上させられることを示しています。