[논문 리뷰] Jasper and Stella: distillation of SOTA embedding models

Donglei Zhang, FulongWang외 2명|arXiv (Cornell University)|2024. 12. 26.

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 교사를 작은 학생 모델로 지식을 전달하는 다중 교사(distillation) 프레임워크를 제안한다.
• 감독 데이터 없이 큰 교사 벡터를 결합한 것을 축소하는 차원 축소 기술을 도입한다.
• 이미지와 텍스트를 모두 인코딩할 수 있는 모델을 훈련시켜 멀티모달 정렬을 시연한다.

제안 방법

• 지식 증류를 사용하여 한 명 또는 다수의 교사 벡터를 모방하도록 학생을 학습시키며 세 가지 손실: cosine_loss, similarity_loss, 그리고 triplet_loss를 사용한다.
• 여러 교사 벡터를 결합하여 고차원 목표를 형성하고 단계적 증류(Stage 1 & 2)를 통해 학생을 학습시킨다.
• 유사성 및 트리플렛 손실로 안내되는 추가 완전히 연결된 층을 사용한 차원 축소 단계(Stage 3)를 적용하여 벡터 크기를 줄인다.
• Stage 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단계 프로세스를 통해 감독 없이 학습된다.
• Stage4는 이미지 자막을 텍스트 임베딩과 정렬하여 이미지–텍스트를 지원하는 인코더를 가능하게 하는 멀티모달 학습을 시연한다.
• 전반적으로 이 접근법은 작은 모델이 증류 및 차원 축소를 통해 강력한 결과를 달성할 수 있음을 보여주며, 실용적인 1.9B 파라미터의 Jasper/Stella 시스템을 제시한다.