Matryoshka Representation Learning

plm(pretrained-language model)을 활용한 임베딩 생성 을 기반으로 벡터검색 을 수행할 때, 일반적으로 큰 차원(768-dim 이상)으로 인해 벡터 검색 서빙에 bottle-neck으로 작용한다 . matryoshka representation learning은 이러한 문제점을 해결하여 성능저하 없이 저차원의 벡터를 활용할 수 있게 하고 , 기 법 을 적용하는 것만 으로도 임베딩 품질 이 개선되는 효과 가 있다 .

• Introduction to Matryoshka Representation Learning: https://huggingface.co/blog/matryoshka

01. MTEB Leaderboard

 Massive Text Embedding Benchkmark는 text embedding을 목표로 하는 다양한 task들을 대상으로 기록되는 model  leaderboard입니다. huggingface를 통해 제공되고 있습니다.

• https://huggingface.co/spaces/mteb/leaderboard

 

 아주 다양한 항목들을 기준으로 갱신되고 있는데, 다음과 같습니다.

• Model Size (Million Parameters)
• Memory Usage (GB, fp32)
• Embedding Dimensions
• Max Tokens
• Average (56 datasets)
• Classification Average (12 datasets)
• Clustering Average (11 datasets)
• PairClassification Average (3 datasets)
• Reranking Average (4 datasets)
• Retrieval Average (15 datasets)
• STS Average (10 datasets)

 많은 semantic search usage, 특히 최근들어 매우 핫해진 RAG(Retrieval Augmented Generation) 분야에서 위 기준들을 근거로하는 sota에 근접한 모델들을 많이 활용하고 있습니다. 그런데 대부분의 모델은 매우 고수준의 임베딩 차원을 가 지고 있습니다.

 

 예를들어 2024.12.19 당시 1위를 기록한 NV-Embed-V2 모델은 4096차원의 임베딩 벡터를 생성합니다. vector search를 위한 index 구축 및 서빙성능 을 고려해보면 다음과 같습니다(검색속도 faiss 기준).

• 4096 차원
  • 각 벡터당 4096 * 4 bytes (float32 기준) = 16,384 bytes (~16KB)
  • 100만개 기준 memory usage: ~15.6GB
  • ~10-20ms per query
• 64 차원
  • 각 벡터당 64 * 4 bytes = 256 bytes
  • 100만개 기준 memory usage:  ~244MB
  • ~1-2ms per query

 일반적으로 임베딩 벡터의 차원이 높을수록 더 좋은 표현력을 가 지고 있지만, 실 서빙단계에서는 컴퓨팅 자원의 제약이 클수도 있음을 알 수 있습니다. matryoshka represent learning은 표현력을 보존하면서 벡터차원을 절감하는 효과가 있어 이런 어려움을 해결하는데 도움을 줄 수 있습니다.

 

02. Matryoshka Representation Learning

 

 training method는 매우 간단하게 이루어졌는데, d차원의 임베딩이 있 을 때, d/2, d/4, d/8, … 차원의 임베딩도 같은 task를 수행합니다. 예를들면, 768차원의 임베딩을 8, 16, 32, 64차원의 임베딩으로 차원을 압축/축소해서 사용하는 것이 아니라 임베딩차원이 낮아지면서 계속해서 포함관계를 가진다. 이러한 특성으로 인해 러시아 인형인 마트료시카 인형과 비슷하다고 하여 마트료시카 표현학습이라는 이름을 붙였습니다.

 

 

정성적 평가

 

 논문에서는 vision task를 기준으로 결과를 설명합니다. 일정 수준의 벡터차원까지는 임베딩 벡터의 성능저하가 크지 않음 을 확인할 수 있으므로 차원축소 관점에서 적절한 차원을 구성하여 Speed & Accuracy Trade-off의 균형을 잡는 것이 중요합니다. 또한, 논문에 따르면 MRL이 원래 임베딩에 비해 성능이 하락하는 수준이 기존보다 훨씬 적으며, 구성도 더 간편합니다.

 

정량적 평가

 

MRL 훈련을 진행한 임베딩이 모든 차원에서 더 좋은 성능을 나타냅니다.

 

낮은 차원에서 기존방법론 대비 더 좋은 성능을 보일뿐 아니라, 768차원 원래 성능과의 비교에서도 더 성능이 더 개선되었음을 알 수 있습니다. MRL의 원래 목적은 차원축소의 자유도의 증가인데, 취지를 넘어 원 임베딩 표현자체도 성능이 개선되었음을 알 수 있습니다.