[Paper Review] Jasper and Stella: distillation of SOTA embedding models

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

Simulation Techniques and Applications8 citations

TL;DR

The paper presents a distillation and dimension-reduction pipeline to produce smaller, efficient text embedding models (Jasper and Stella) that achieve strong MTEB scores, including multimodal alignment with images, without requiring supervised data.

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).

Motivation & Objective

• Motivate reducing latency and storage for dense retrieval by distilling large embedding models.
• Propose a multi-teacher distillation framework to transfer knowledge from several SOTA teachers to a smaller student.
• Introduce a dimension-reduction technique to shrink large teacher-combined vectors without supervised data.
• Demonstrate multimodal alignment by training a model that can encode both images and text.

Proposed method

• Use knowledge distillation to train a student to imitate one or more teacher vectors with three losses: cosine_loss, similarity_loss, and triplet_loss.
• Concatenate multiple teacher vectors to form high-dimensional targets and train the student via stage-wise distillation (Stage 1 & 2).
• Apply a dimension-reduction stage (Stage 3) using additional fully connected layers guided by similarity and triplet losses to reduce vector size.
• Stage 4 trains a visual encoder to align image captions with text embeddings, enabling image–text multimodal encoding using the same loss framework.
• The training setup uses unsupervised data; no supervised labels are required.

Figure 1: The model architecture of Jasper model

Experimental results

Research questions

• RQ1 Can distillation from multiple SOTA embedding teachers yield a smaller student model with competitive performance?
• RQ2 Does a dedicated dimension-reduction stage preserve performance when compressing concatenated teacher vectors?
• RQ3 Can a multimodal encoder be trained to align visual and textual representations without supervised data?
• RQ4 What are the practical performance characteristics (MTEB scores) of Jasper and Stella compared to existing models under similar parameter budgets?

Key findings

• Jasper and Stella achieve strong MTEB results, with Jasper averaging 72.02 on 56 datasets, outperforming several baselines under 2B parameters.
• NV-Embed-v2 remains a stronger baseline at 72.31 avg, while Jasper is competitive across classifications, clustering, pair classification, reranking, retrieval, STS, and summarization tasks.
• The models are trained via a four-stage process including distillation from multiple teachers, dimensionality reduction, and multimodal alignment, all using unsupervised data.
• Stage4 demonstrates multimodal training by aligning image captions with text embeddings, enabling an image–text capable encoder.
• Overall, the approach shows that small models can achieve strong results through distillation and dimension reduction, with a practical 1.9B parameter Jasper/Stella system.