Source-linked AI summary
Jasper and Stella: distillation of SOTA embedding models
Dun Zhang, Jiacheng Li, Ziyang Zeng, Fulong Wang
TL;DR
High-performing embedding models can be costly to deploy because of their parameter counts and vector dimensions. The paper distills multiple teachers into a smaller student using three losses and MRL, with Jasper reaching No.3 on MTEB at the 2B scale. Multimodal alignment remains preliminary and constrained by limited training resources.
Problem
High-performing embedding models often have many parameters and high-dimensional vectors, creating slow inference and high storage costs.
Method
A multi-stage framework distills multiple larger teachers into a smaller student with three losses, while MRL reduces vector dimensionality.
Results
Jasper achieved the No.3 MTEB position with a 71.54 average across 56 datasets at 2B parameters.
Takeaways & Limitations
Jasper delivers state-of-the-art performance at the 2B parameter scale and results comparable to top-ranked 7B embedding models.
Takeaways & Limitations
The multimodal training provides only preliminary text-image alignment, with limited image-encoding capability and oscillatory stage-4 loss behavior.
Abstract
from arXiv · showhide
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).
1 Introduction
High-performing embedding models often require substantial parameters and vector dimensions, motivating a multi-stage distillation framework that produces a smaller student model. Jasper combines multiple-teacher distillation with dimensionality reduction and achieves strong MTEB performance at 2B parameters.
- Motivation: High-performing embedding models often use many parameters and high-dimensional vectors, causing slow inference and high storage costs.NV-Embed-v2 and bge-en-icl are cited as 7B-parameter models with 4096-dimensional representations.
- Method: The proposed framework distills knowledge from multiple larger teacher embedding models into a smaller student through three loss functions.The losses progress from representation matching to semantic similarity and relative ranking information.
- Method: MRL-based training compresses the student’s vector representation, while a vision encoder and self-distillation align visual and textual embeddings.The overall process uses a four-stage distillation approach.
- Results: 71.54 average across 56 datasets is reported for Jasper, a 2B-parameter student model on MTEB.Jasper is described as comparable to 7B models and significantly better than models with fewer than 2B parameters.
- Results: Jasper achieved the No.3 position on the MTEB leaderboard as of December 24, 2024.The paper compares its results with top-ranked 7B embedding models and smaller models.
2 Methods
Jasper combines language and vision encoders with multi-teacher distillation, three progressively broader losses, and MRL-based dimensionality reduction. The framework also explores self-distillation for producing shorter embeddings and preliminary text-image alignment.
- 2.2 Model Architecture: Jasper uses a language model, vision encoder, pooler, and fully connected layers to produce final embeddings.The language model uses mean pooling, while the vision pathway maps image tokens to the language embedding space before projection.
- 2.3 Stage 1&2: Distillation from Multiple Teachers: Multiple teacher vectors are concatenated, requiring the student representation to match their combined 12,288-dimensional size during early distillation.The teachers are NV-Embed-v2 with 4,096 dimensions and stella_en_1.5B_v5 with 8,192 dimensions.
- 2.3 Stage 1&2: Distillation from Multiple Teachers: Three losses progress from vector alignment to pairwise similarity and relative ranking, with the final objective combining them using weighted coefficients.Cosine loss aligns absolute representations, similarity loss compares semantic matching, and relative similarity distillation uses positive-negative pair preferences.
- 2.4 Stage 3: Dimension Reduction: MRL introduces three independent fully connected layers that generate lower-dimensional vectors, while FC1 preserves the full teacher-aligned representation.The reduced vectors use similarity and relative-similarity losses, whereas FC1 continues using all three losses.
- 2.4 Stage 3: Dimension Reduction: Self-distillation uses FC1’s 12,288-dimensional outputs as teachers for shorter vectors, but the paper does not experimentally evaluate this approach’s specific merits.This method is intended to reduce embedding dimensionality using only unsupervised data and the model itself.
- 2.5 Stage 4: Unlock Multimodal Potential: Stage 4 trains only the visual encoder on image-caption pairs, using caption representations as teachers and image representations as students.The resulting text-image alignment is preliminary and leaves significant room for improvement.
3 Experiments
The experiments evaluate Jasper and Stella across the full MTEB benchmark and compare Jasper with frontier models and smaller models. Jasper significantly outperforms models below 2B parameters while matching results from 7B models.
- 3.1 Implementation details: Jasper contains 1.9B parameters and is initialized from stella_en_1.5B_v5 and SigLIP, with NV-Embed-v2 and Stella serving as teacher models.The reported training setup uses four stages with specified checkpoints, batch sizes, and learning rates.
- 3.1 Implementation details: The training data combines 80% FineWeb-Edu with 20% sentence-transformers embedding-training data, plus BAAI/Infinity-MM caption data for vision training.The additional text dataset supplies questions alongside passages, while caption data is used in stage 4.
- 3 Experiments: Jasper is evaluated on the full MTEB benchmark, covering retrieval, reranking, classification, clustering, pair classification, semantic textual similarity, and summarization.The benchmark includes 56 datasets across seven task categories.
- 3.3 Results: Jasper significantly outperforms models with fewer than 2B parameters across the overall MTEB score and seven subcategory averages.Table 1 reports average scores for overall performance and seven MTEB task categories.
- 3.3 Results: Jasper produces results comparable to models with 7B parameters despite having only 2B-scale capacity.This comparison is reported as part of the MTEB evaluation.
4 Discussion
Jasper is evaluated as an instruction-based embedding model and is reported to be robust to varied prompts. Its multimodal capability remains limited because visual encoding received only basic training under resource constraints.
- Instruction robustness: Jasper is an instruction-based embedding model evaluated with original prompts and GPT-4o-generated instruction variants.The experiment compares Jasper’s performance under similar instructions.
- Instruction robustness: Jasper is reported to be robust to different instructions and able to understand them accurately.This conclusion is based on short evaluation tasks using modified instructions.
- Vision encoding: Basic image encoding was implemented because time and resource constraints limited multimodal training.The planned visual-language alignment and potential VQA-based contrastive-learning stages were not fully developed.
- Vision encoding: Oscillatory loss behavior during stage 4 indicates considerable room for improving multimodal training.The paper identifies further multimodal alignment as future work.
5 Conclusion
The paper presents Jasper’s distillation-based training procedure, combining multiple teacher models, three loss functions, and MRL-based dimensionality reduction. Experiments on MTEB report state-of-the-art performance at the 2B parameter scale and results comparable to top-ranked 7B models.
- Conclusion: The paper presents a distillation-based training procedure for the Jasper model.The conclusion introduces the paper’s central training contribution.
- Conclusion: Three loss functions distill multiple large teacher embedding models into a student model from diverse perspectives.The conclusion describes this as a central component of the training procedure.
- Conclusion: MRL-based training reduces the vector dimensionality of the student model.This dimensionality reduction follows the multi-teacher distillation procedure.
- Conclusion: Jasper achieves state-of-the-art performance at the 2B parameter scale on MTEB.The conclusion also reports comparable results to other top-ranked embedding models with 7B parameters.
- Conclusion: Future work will further explore alignment between multiple modalities.The conclusion identifies multimodal alignment as an area for continued investigation.