Jasper and Stella: distillation of SOTA embedding models

在许多深度学习应用中，如常见问题解答（FAQ）和检索增强生成（RAG），密集检索是一个关键组件。在这个过程中，嵌入模型将原始文本转换为数值向量。然而，当前在文本嵌入基准测试（如大规模文本嵌入基准测试MTEB）中表现出色的嵌入模型通常具有大量参数和高维向量，这给其在实际场景中的应用带来了挑战。为了解决这一问题，我们提出了一种新颖的多阶段蒸馏框架，使较小的学生嵌入模型能够通过三种精心设计的损失函数从多个较大的教师嵌入模型中进行蒸馏。同时，我们利用套娃表示学习（MRL）有效降低了学生嵌入模型的向量维度。我们的学生模型Jasper基于Stella嵌入模型，拥有20亿个参数，在MTEB排行榜上获得了第三名（截至2024年12月24日），在56个数据集上的平均得分为71.54。我们已在Hugging Face Hub发布了该模型和数据（https://huggingface.co/infgrad/jasper_en_vision_language_v1）（https://huggingface.co/datasets/infgrad/jasper_text_distill_dataset），训练代码也已在项目仓库中提供（https://github.com/NLPJCL/RAG-Retrieval）。

该论文旨在解决当前在密集检索应用中，如FAQ和RAG，所使用的高质量文本嵌入模型参数过多、向量维度过高所带来的实际应用挑战。这是一个现有问题，但目前缺乏有效的解决方案来平衡模型性能与资源消耗。

关键思路在于提出一个多阶段蒸馏框架，使较小的学生模型能够从多个较大的教师模型中学习，并通过三种精心设计的损失函数进行优化。同时引入了Matryoshka表示学习（MRL）以有效降低学生模型的向量维度。这一方法不仅减少了模型的复杂度，还保持了较高的性能水平。

实验设计包括将新模型Jasper应用于56个不同的数据集，并在MTEB排行榜上取得了第三名的成绩，平均得分为71.54。此外，作者团队开源了模型、训练代码及数据集，为后续研究提供了宝贵资源。这些工作为进一步探索更小、更快且高效的嵌入模型铺平了道路。

近期相关研究包括《Dense Passage Retrieval for Open-Domain Question Answering》和《Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks》，这些研究同样关注于提高文本嵌入的质量和效率。此外，《Efficient Estimation of Word Representations in Vector Space》也探讨了如何减少模型参数而不牺牲性能。