Matryoshka 임베딩 모델 소개

Matryoshka 임베딩 모델은 성능에 큰 손실 없이 더 작은 차원으로 축소할 수 있는 임베딩을 생성할 수 있게 합니다. 이 접근 방식은 실무자가 저장 비용, 처리 속도 및 정확도 요구 사항에 따라 임베딩 솔루션을 동적으로 확장할 수 있도록 합니다.

Matryoshka 임베딩 작동 방식

Matryoshka 임베딩은 Matryoshka Representation Learning (MRL)에 기반한 기술로, 모델이 임베딩 벡터의 가장 초기 차원에 가장 중요한 정보를 저장하도록 장려합니다. 러시아의 중첩 인형처럼, 모델은 축소된 임베딩 버전이 여전히 다운스트림 작업에 유용할 만큼 충분한 정보를 유지하도록 학습됩니다.

적용 분야 및 장점

가변 크기 임베딩은 AI 실무자에게 두 가지 주요 이점을 제공합니다:

  1. Shortlisting and Reranking: 사용자는 작은 축소된 임베딩을 사용하여 초기의 매우 효율적인 "shortlist" 검색을 수행할 수 있습니다. 후보 집합이 식별되면 전체 차원의 임베딩을 사용하여 결과를 재정렬해 높은 정밀도를 얻을 수 있습니다.
  2. Resource Optimization: Matryoshka 모델은 배포 환경의 특정 제약에 따라 저장 비용, 처리 속도 및 성능 사이의 맞춤형 트레이드오프를 가능하게 합니다.

Matryoshka 임베딩 모델 훈련

이론적 프레임워크

표준 임베딩 훈련에서는 전체 크기의 임베딩에 손실 함수를 적용합니다. 반면 Matryoshka 훈련은 전체 크기 임베딩뿐만 아니라 다양한 축소 차원(예: 768, 512, 256, 128, 64)의 임베딩에 대한 손실 값을 계산합니다. 이러한 손실을 합산하여 최종 손실 값을 만들고, 옵티마이저가 모델 가중치를 조정하여 가장 중요한 정보를 "앞쪽에" 배치하도록 합니다.

Sentence Transformers에서 구현

Sentence Transformers는 MatryoshkaLoss를 통해 Matryoshka 모델 지원을 구현했습니다. 이 래퍼 손실 함수는 CoSENTLoss와 같은 기본 손실을 임베딩의 여러 축소 부분에 적용합니다. MatryoshkaLoss를 사용한 훈련은 훈련 시간에 큰 오버헤드를 발생시키지 않습니다.

Matryoshka 임베딩 모델 사용

일반 추론

추론은 표준 임베딩 모델과 유사하게 작동합니다. 유일한 차이점은 임베딩이 생성된 후 선택적으로 더 작은 차원으로 축소될 수 있다는 점입니다. 임베딩이 원래 정규화되었다면, 축소 후 일관성을 유지하기 위해 다시 정규화해야 합니다.

Sentence Transformers에서 구현

SentenceTransformer 클래스를 사용하면, 모델 초기화 시 truncate_dim 인자를 지정할 수 있습니다. encode 함수는 지정된 크기로 출력 임베딩을 자동으로 축소합니다.

nomic-ai/nomic-embed-text-v1.5와 같은 특정 모델의 경우, 임베딩을 축소하기 전에 F.layer_norm을 적용하는 맞춤형 아키텍처가 필요합니다.

성능 결과

STSBenchmark 테스트 세트에서 Matryoshka 모델(tomaarsen/mpnet-base-nli-matryoshka)과 표준 모델(tomaarsen/mpnet-base-nli)을 비교한 실험 결과는 다음과 같습니다:

  • Higher Accuracy: Matryoshka 모델은 모든 테스트 차원에서 더 높은 Spearman 유사도를 달성했습니다.
  • Performance Retention: Matryoshka 모델의 성능 저하 속도는 표준 모델보다 훨씬 느립니다. 원래 크기의 8.3%로 축소했을 때에도 Matryoshka 모델은 최대 성능의 98.37%를 유지했으며, 표준 모델은 96.46%에 머물렀습니다.

이 결과는 Matryoshka 임베딩이 성능에 큰 영향을 주지 않으면서 저장 공간을 크게 줄이고 검색 속도를 높일 수 있음을 보여줍니다.

Sources