Kusupati et al.이 Matryoshka Representation Learning 방식으로 임베딩 모델을 학습시켜 원본 크기의 8.3%만 사용하면서 98.37% 성능 유지

Context

최신 임베딩 모델들은 성능 향상을 위해 점점 높은 차원(dimension)을 사용하고 있으나, 이는 검색이나 분류 같은 다운스트림 작업의 저장소 비용과 처리 속도를 증가시킨다.

Technical Solution

• Matryoshka 학습 방식 도입: 768, 512, 256, 128, 64 등 다양한 차원에서 동시에 손실함수를 적용하여 중요 정보를 임베딩 앞부분에 집중시킴
• Sentence Transformers 프레임워크의 MatryoshkaLoss 활용: CoSENTLoss 같은 기존 손실함수를 여러 차원에서 병렬로 계산하고 가중치를 적용하여 합산
• 차원 축소 기반 2단계 검색 전략: 작은 차원으로 빠르게 후보를 shortlist한 후 남은 임베딩을 전체 차원으로 재처리
• 임베딩 앞부분에 정보 집중: 모델이 원본 임베딩의 초반 치수(dimension)에 가장 중요한 정보를 배치하도록 학습되어 절단(truncation) 시에도 성능 손실 최소화

Impact

• STSBenchmark 테스트 셋에서 원본 임베딩 크기의 8.3% 사용 시 Matryoshka 모델은 98.37% 성능 유지, 표준 모델은 96.46% 성능만 유지
• Matryoshka 모델이 모든 차원에서 표준 모델보다 높은 Spearman 유사도 달성
• 학습 시간에 측정 가능한 오버헤드 없음

Key Takeaway

임베딩 차원을 가변적으로 설정 가능한 모델 설계는 저장소 비용, 처리 속도, 성능 간의 선택지를 제공하므로, 다양한 배포 환경에서 실시간 조정과 최적화를 가능하게 한다.