[論文レビュー] Jasper and Stella: distillation of SOTA embedding models

Donglei Zhang, FulongWang+2 more|arXiv (Cornell University)|Dec 26, 2024

Simulation Techniques and Applications被引用数 8

ひとこと要約

本論文は、JasperとStellaという小型で効率的なテキスト埋め込みモデルを生み出す蒸留と次元削減のパイプラインを提示し、監督データを必要とせずにマルチモーダル整合性を含む強力なMTEBスコアを達成する。

ABSTRACT

A crucial component in many deep learning applications, such as Frequently Asked Questions (FAQ) and Retrieval-Augmented Generation (RAG), is dense retrieval. In this process, embedding models transform raw text into numerical vectors. However, the embedding models that currently excel on text embedding benchmarks, like the Massive Text Embedding Benchmark (MTEB), often have numerous parameters and high vector dimensionality. This poses challenges for their application in real-world scenarios. To address this issue, we propose a novel multi-stage distillation framework that enables a smaller student embedding model to distill multiple larger teacher embedding models through three carefully designed losses. Meanwhile, we utilize Matryoshka Representation Learning (MRL) to reduce the vector dimensionality of the student embedding model effectively. Our student model named Jasper with 2 billion parameters, built upon the Stella embedding model, obtained the No.3 position on the MTEB leaderboard (as of December 24, 2024), achieving an average 71.54 score across 56 datasets. We have released the model and data on the Hugging Face Hub (https://huggingface.co/infgrad/jasper_en_vision_language_v1) (https://huggingface.co/datasets/infgrad/jasper_text_distill_dataset), and the training codes are available in this project repository (https://github.com/NLPJCL/RAG-Retrieval).

研究の動機と目的

• dense retrievalの待機時間とストレージを削減するための大規模埋め込みモデルの蒸留による動機付け。
• 複数のSOTA教師から小さな生徒モデルへ知識を移 transferする多教師蒸留フレームワークを提案。
• 監督データなしで大きな教師結合ベクトルを縮小する次元削減技法を導入。
• 画像とテキストの両方をエンコードできるマルチモーダル整合性を訓練で実現。

提案手法

• 知識蒸留を用いて、1つ以上の教師ベクトルを模倣する生徒を訓練し、3つの損失関数（cosine_loss、similarity_loss、triplet_loss）を適用。
• 複数の教師ベクトルを結合して高次元のターゲットを形成し、Stage 1とStage 2の段階的蒸留で生徒を訓練。
• similarityとtriplet損失に guided された追加の全結合層を用いた次元削減段階（Stage 3）でベクトルサイズを縮小。
• Stage 4では視覚的エンコーダを訓練し、画像キャプションをテキスト埋め込みと整列させ、同一の損失フレームワークを用いて画像-テキストのマルチモーダルエンコーダを可能にする。
• 訓練設定は非教師データを使用；監督ラベルは不要。

Figure 1: The model architecture of Jasper model

実験結果

リサーチクエスチョン

• RQ1 複数のSOTA埋め込み教師からの蒸留は、競争力のある性能を持つ小型の生徒モデルを生み出せるか？
• RQ2 結合された教師ベクトルを圧縮する際、専用の次元削減段は性能を保持できるか？
• RQ3 監督データなしで視覚的・ textual 表現を整列させるマルチモーダルエンコーダを訓練できるか？
• RQ4 JasperとStellaの実用的な性能特性（MTEBスコア）は、同等のパラメータ予算の既存モデルと比べてどうか？

主な発見

モデル	平均値（56データセット）	分類	クラスタリング	ペア分類	リランキング	検索	STS	要約
NV-Embed-v2	72.31	90.37	58.46	88.67	60.65	62.65	84.31	30.7
jasper(our model)	72.02	88.49	58.04	88.07	60.91	63.12	84.67	31.42
bge-en-icl	71.67	88.95	57.89	88.14	59.86	62.16	84.24	30.77
stella_en_1.5B_v5	71.19	87.63	57.69	88.07	61.21	61.01	84.51	31.49

• JasperとStellaは強力なMTEB結果を達成し、Jasperは56データセットで平均72.02で、2Bパラメータ未満のベースラインを上回る。
• NV-Embed-v2は平均72.31で依然として強力なベースラインであり、Jasperは分類、クラスタリング、ペア分類、リランキング、検索、STS、要約タスクで競合的。
• 本モデルは複数の教師からの蒸留、次元削減、マルチモーダル整合性を含む4段階のプロセスで訓練され、すべて非教師データを使用。
• Stage4は画像キャプションをテキスト埋め込みと整列させるマルチモーダル訓練を実証し、画像-テキスト対応エンコーダを実現。
• 全体として、本アプローチは蒸留と次元削減を通じて小さなモデルでも強力な結果を達成できることを示しており、実用的な1.9BパラメータのJasper/Stellaシステムを提示。