[论文解读] 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).

研究动机与目标

• 通过蒸馏大嵌入模型来降低密集检索的延迟和存储需求。
• 提出一个多教师蒸馏框架，将来自若干SOTA教师的知识转移给更小的学生模型。
• 引入一个降维技术，在无监督数据的情况下缩小大型教师-组合向量。
• 通过训练一个能够同时编码图像与文本的模型来展示多模态对齐。

提出的方法

• 使用知识蒸馏让学生模仿一个或多个教师向量，使用三种损失：cosine_loss、similarity_loss 和 triplet_loss。
• 将多个教师向量拼接成高维目标，通过阶段性蒸馏（阶段1与阶段2）来训练学生。
• 应用一个降维阶段（阶段3），通过额外的全连接层并在相似性和三元组损失引导下降低向量大小。
• 阶段4训练一个视觉编码器，使图像标题/描述与文本嵌入对齐，从而在同一损失框架下实现图像–文本多模态编码。
• 训练设置使用无监督数据；不需要监督标签。

Figure 1: The model architecture of Jasper model

实验结果

研究问题

• RQ1 从多位SOTA嵌入教师进行蒸馏是否能得到具有竞争性能的更小学生模型？
• RQ2 在压缩拼接的教师向量时，专用的降维阶段是否能保留性能？
• RQ3 能否训练出一个多模态编码器，在不使用监督数据的情况下对齐视觉与文本表示？
• RQ4 在类似参数预算下，Jasper 与 Stella 相比现有模型的实际性能特征（MTEB 分数）如何？

主要发现

• Jasper 与 Stella 实现了强劲的 MTEB 结果，Jasper 在56个数据集上平均 72.02，超过若干低于2B参数的基线。
• NV-Embed-v2 仍然在平均分 72.31 时表现更强；而 Jasper 在分类、聚类、对成对分类、重新排序、检索、STS 与摘要等任务上具竞争力。
• 模型通过包括来自多位教师的蒸馏、降维和多模态对齐在内的四阶段流程训练，且全部使用无监督数据。
• 阶段4 通过将图像标题与文本嵌入对齐来演示多模态训练，使得编码器具备图像–文本能力。
• 总体而言，该方法显示小模型可通过蒸馏与降维实现强结果，建立了一个实际的 1.9B 参数 Jasper/Stella 系统。