마트료시카 표현 학습 설명: OpenAI의 효율적인 텍스트 임베딩 이면의 방법
실제 머신 러닝 모델을 개발할 때는 항상 비용과 성능 사이의 트레이드오프가 수반됩니다. 예를 들어, 학습 중에 사용하는 모델과 데이터셋이 클수록, 최종적으로 학습된 모델의 성능은 더 뛰어나게 됩니다. 하지만 이러한 향상된 성능에는 더 긴 시간과 더 높은 계산 비용이 따르며, 학습에 훨씬 더 긴 시간이 필요합니다. 추론 중에도 동일한 원리가 적용되는데, 더 큰 모델은 저장을 위해 더 많은 메모리를 필요로 하는 더 큰 특징 표현을 생성하는 경향이 있습니다.
머신 러닝 학습과 추론 중의 요구사항은 사용 사례마다 크게 달라질 수 있으므로, 비용 절감을 위해 모델 성능의 일부를 약간 희생할 수 있게 해주는 방법이 필수적입니다. 바로 이 지점에서 Matryoshka Representation Learning (MRL) 과 같은 방법이 등장합니다. 예를 들어, OpenAI의 text-embedding-3-small 모델은 MRL을 활용하여 개발자가 핵심 개념 표현 속성을 유지하면서 임베딩을 더 짧게 만들 수 있도록 합니다. 임베딩 차원의 조정을 가능하게 함으로써, MRL은 비용 효율성과 모델 성능 사이의 완벽한 균형을 맞추는 데 도움을 줍니다.
이 글에서는 MRL이 어떻게 작동하는지, 그 구현 방식, 그리고 확장 가능하고 효율적인 머신 러닝 모델을 가능하게 하는 방법을 살펴봅니다. MRL의 동기부터 시작해 보겠습니다.
Matryoshka Representation Learning (MRL)의 동기
비용-성능 트레이드오프는 어떤 머신 러닝 모델을 개발하거나 사용하기 전에 항상 고려해야 할 사항입니다.
학습 중에는 모델의 파라미터와 학습 데이터의 양이 최종 성능에 직접적인 영향을 미칩니다. 더 큰 모델과 학습 데이터셋은 더 뛰어난 모델로 이어지는 경향이 있습니다. 하지만 이러한 모델을 학습시키는 데 드는 계산 비용도 더 높습니다.
한편, 추론 중에는 더 큰 모델이 더 큰 특징 표현을 생성합니다. 이러한 특징 표현의 증가는 이를 저장하기 위해 더 많은 메모리가 필요하다는 뜻입니다.
정보 검색의 맥락에서 BERT를 예로 들어보겠습니다. BERT는 텍스트 분류, Named Entity Recognition (NER), 질의응답(QnA) 등 여러 벤치마크 데이터셋에서 최첨단 성능을 달성한 가장 초기의 Transformer 기반 딥러닝 모델 중 하나입니다.
그림: BERT base 모델을 사용한 임베딩 생성 워크플로.
간단히 말해, BERT는 문장이나 단어를 입력으로 받아 고정 크기의 embedding으로 변환합니다. 이 임베딩은 원래 입력의 의미적 의미를 포착하며, 그 크기는 사용된 특정 모델 변형에 따라 달라집니다.
BERT에는 두 가지 주요 변형이 있습니다: base 모델과 large 모델입니다. base 모델은 768차원 임베딩을 생성하는 반면, large 모델은 1024차원 임베딩을 생성합니다.
정보 검색 사용 사례를 위해 vector database에 1,000만 개의 임베딩을 저장한다고 가정해 보겠습니다. BERT base 모델을 FP32 형식으로 사용한다면, 이러한 임베딩을 저장하는 데 대략 768×10M×4=30.72 GB의 메모리가 필요합니다. 반면, 동일한 시나리오에서 BERT large 모델을 사용하면 약 40.9 GB가 필요합니다. 또한 임베딩 크기가 커지면 similarity searches의 계산 복잡도가 증가하여 검색 프로세스가 느려질 수 있습니다. 그러나 더 큰 임베딩을 사용하는 주된 장점은 더 작은 임베딩에 비해 관련성이 높은 정보를 검색하는 능력이 향상된다는 점입니다.
정보 검색에서 이러한 문제들에 대한 이상적인 해결책 중 하나는 전체 프로세스를 두 부분으로 나누는 것입니다: 후보 선별(shortlisting)과 재순위화(reranking). 후보 선별은 Milvus와 같은 벡터 데이터베이스의 방대한 컬렉션에서 초기 후보 문서 세트를 검색하는 것을 의미합니다. 그런 다음 재순위화는 이렇게 선별된 후보들을 가져와 최종 결과의 관련성을 극대화하도록 다시 정렬합니다.
그림- 후보 선별 및 재순위화 워크플로
그림: 후보 선별 및 재순위화 워크플로.
후보 선별을 수행하기 위해 알고리즘은 쿼리 임베딩과 벡터 데이터베이스에 저장된 대규모 임베딩 컬렉션 간의 유사도를 계산합니다. 따라서 더 작은 임베딩 크기를 사용하면 이 계산을 더 효율적이고 빠르게 만들 수 있습니다. 반대로 재순위화는 정밀도를 강조하므로, 더 큰 임베딩 크기를 사용하면 이점을 얻을 수 있습니다.
이 방법의 문제는 특정 모델을 선택하면 해당 모델이 생성하는 임베딩 또는 특징 표현의 크기가 고정된다는 점입니다. 이러한 유연성 부족은 단일 모델에서 서로 다른 크기의 임베딩을 사용하는 능력을 제한합니다.
그림- 고정 크기 임베딩을 생성하는 모델의 워크플로..png
그림: 고정 크기 임베딩을 생성하는 모델의 워크플로.
Matryoshka Representation Learning (MRL) 접근법은 이 문제에 대한 흥미로운 해결책을 제공합니다: 이 방법으로 학습된 모델은 다양한 크기의 임베딩을 생성할 수 있습니다. 다음 섹션에서 이 접근법의 메커니즘을 자세히 살펴보겠습니다.
Matryoshka Representation Learning (MRL)이란 무엇인가요?
Matryoshka Representation Learning (MRL)은 단일 모델 내에서 다중 스케일 표현을 생성하도록 신경망을 학습시키는 방법입니다. 작은 인형들이 더 큰 인형 안에 들어가는 마트료시카 인형에서 영감을 받은 MRL은 단일 순전파를 사용해 모델이 다양한 크기(거친 수준에서 세밀한 수준까지)의 표현을 출력할 수 있게 합니다. 이 접근법은 딥러닝 모델이 데이터의 복잡한 관계와 미묘한 차이를 더 잘 이해하도록 해줍니다. MRL은 특히 의미론적 검색, 정보 검색, 다국어 처리, 그리고 서로 다른 추상화 수준에서 데이터의 섬세한 표현이 필요한 모든 애플리케이션과 같은 작업에 효과적입니다.
여러 세부 레이어를 가진 Matryoshka 임베딩 시각화
그림: 여러 세부 레이어를 가진 Matryoshka 임베딩 시각화
Matryoshka Representation Learning (MRL) 접근법을 채택한 인기 임베딩 모델로는 OpenAI의 text-embedding-3-large, Nomic의 nomic-embed-text-v1, Alibaba의 gte-multilingual-base가 있습니다.
MRL 접근법의 작동 방식
MRL 접근법은 어떤 머신러닝 모델에서든 여러 크기의 특징 표현을 추출할 수 있게 해줍니다. 예를 들어, 임베딩의 원래 1024차원을 사용하는 대신 처음 16, 32, 64, 128 또는 256차원(또는 우리가 선택하는 어떤 차원이든)을 활용할 수 있습니다. 이러한 기능의 핵심은 MRL 접근법을 구현할 때 모델이 학습되는 방식에 있습니다.
그림- MRL 손실 함수 학습과 추론 중 사용 사례
그림: MRL 손실 함수 학습과 추론 중 사용 사례. 출처.
MRL을 사용한 모델 학습 중에는 표준 모델 학습에서 일반적인 것처럼 하나의 손실 함수만 최적화하는 것이 아니라 여러 손실 함수를 최적화합니다. 각 손실 함수는 사전에 정의한 특정 차원에 대해 특징 표현을 최적화하는 것을 목표로 합니다. 다섯 가지 서로 다른 특징 차원을 설정하면, 학습 중에 다섯 개의 서로 다른 손실 함수를 최적화해야 합니다.
MRL 접근 방식을 사용해 BERT large 모델을 학습시키고 싶다고 해봅시다. 이전 섹션에서 언급했듯이, 이 모델은 크기 1024의 임베딩을 생성합니다. 특징 차원을 32, 64, 128, 256, 1024로 설정하면, 모델이 해당 차원에서 최적화된 임베딩을 생성하도록 학습 중에 다섯 개의 서로 다른 손실 함수를 최적화하게 됩니다.
이러한 다양한 차원 전반의 최적화는 간단합니다. MRL은 전체 손실 함수를 각 개별 차원에 대한 손실의 합으로 분해합니다. 위의 예를 사용하면, 이는 다음과 같이 표현할 수 있습니다.
MRL로 학습된 모델의 최적화 과정은 다음과 같이 더 공식적으로 표현할 수 있습니다.
그림: MRL 접근 방식의 최적화 방정식. 출처.
보시다시피, 위 손실 함수의 정의는 상당히 일반적입니다. 따라서 MRL 접근 방식은 거의 모든 모델에 적용할 수 있으며 모델의 아키텍처와 완전히 독립적입니다. 또한 BERT와 같은 사전 학습된 모델이나 다른 transformer 기반 모델을 파인튜닝하여 다양한 크기의 임베딩을 출력하도록 할 수도 있습니다. 위의 시나리오를 이어가면, 이제 학습된 BERT large 모델에서 1024차원 임베딩과 32, 64, 128, 256차원 임베딩을 생성할 수 있습니다.
전반적으로, MRL로 모델을 학습한 후에는 학습된 모델이 생성한 특징의 초기 차원이 이후 차원보다 더 중요한 정보를 담고 있습니다. 처음 몇 개의 차원은 상위 수준의 세부 정보를 포함하고, 이후 차원은 더 세분화된 정보에 초점을 맞추며, 이는 마트료시카 인형의 구조를 모방합니다.
하지만 이것이 더 짧은 임베딩이 항상 더 긴 임베딩을 단순히 잘라낸 버전이라는 의미는 아닙니다. 짧은 임베딩의 각 요소 값은 학습 중 각 특징 차원에 적용된 스케일링 계수 때문에 긴 임베딩의 값과 다를 수 있습니다. 그러나 각 차원의 스케일링 계수를 동일하게 설정하면, 짧은 임베딩과 긴 임베딩 간의 각 요소 값이 동일해 보일 수 있습니다.
MRL 실험 결과
MRL 접근 방식은 텍스트, 비전, 비전-텍스트를 포함한 다양한 모달리티의 머신러닝 모델에서 평가되었습니다. ResNet50 및 ViT 모델은 비전 기반 모델을 나타내고, BERT는 텍스트 기반 모델을 나타내며, ALIGN은 비전과 텍스트의 조합을 나타냅니다. 이러한 모델은 주로 두 가지 일반적인 사용 사례인 분류 및 검색에서 평가되었습니다. 먼저 분류에 대해 논의해 보겠습니다.
ImageNet-1K 데이터셋에서 학습된 ResNet50 모델의 성능을 독립적으로 학습된 표준 ResNet50과 비교했을 때, MRL 모델은 다양한 특징 표현 크기에서 비슷한 top-1 정확도를 달성합니다.
다운스트림 작업에서 특징 표현의 유용성을 더 평가하기 위해, 각 특징 표현 크기에 대한 1-nearest neighbor(1-NN)의 정확도도 측정했습니다. 1-NN을 계산하기 위한 설정은 다음과 같습니다. 데이터베이스에 1.3K개의 이미지 샘플이 주어졌을 때, 작업은 50K개의 쿼리 각각에 대해 가장 가까운 이웃을 찾는 것입니다. 이 설정에서 MRL로 학습된 ResNet50은 각 특징 표현 크기에서 고정 특징 대응 모델보다 최대 2% 더 높은 정확도를 보입니다.
Figure- Top-1 accuracy and 1-NN accuracy of ResNet50 on ImageNet-1K
그림: ImageNet-1K에서 ResNet50의 Top-1 정확도 및 1-NN 정확도. 출처.
한편, JFT-300M 데이터셋에서 MRL로 학습된 ViT 모델의 성능도 모든 표현 크기에서 매우 경쟁력이 있습니다. 1-NN 정확도는 고정 크기 특징 표현으로 학습된 ViT의 정확도와 비슷합니다. 아래 그림에 표시된 것처럼, MRL 모델의 성능은 낮은 특징 표현에서 고정 크기 대응 모델보다도 더 우수한데, 이는 부분적으로 낮은 차원을 표현하기 위해 고정 크기 모델에서 무작위 특징이 선택되기 때문입니다. MRL 접근법을 사용해 학습된 ALIGN 모델에서도 유사한 추세가 관찰되며, 그 성능은 고정 크기 표현으로 학습된 ALIGN 모델의 성능과 일치합니다.
Figure- 1-NN accuracy of ViT and ALIGN on ImageNet-1K
그림: ImageNet-1K에서 ViT 및 ALIGN의 1-NN 정확도. 출처.
MRL 접근법을 사용하는 주요 장점은 동일한 모델에서 서로 다른 크기의 특징 표현 간에 전환할 수 있는 유연성입니다. 이 장점을 충분히 활용하기 위해 ImageNet-1K에서 ResNet50을 사용한 적응형 분류 테스트도 수행했습니다.
이 설정은 더 작은 특징 표현에서 더 큰 특징 표현으로의 전환을 결정하기 위해 최대 softmax 확률을 학습하는 것을 포함합니다. 테스트 결과에 따르면, 37차원 특징 표현을 사용하는 MRL 모델은 고정 512차원 특징 표현으로 학습된 ResNet50 모델만큼 정확하면서도 2048차원 모델보다 정확도가 단 0.8% 낮습니다.
Figure- Top 1 accuracy of ResNet50 model in adaptive classification (left) and mAP of ResNet50 in image retrieval use case (right)
그림: 적응형 분류에서 ResNet50 모델의 Top 1 정확도(왼쪽)와 이미지 검색 사용 사례에서 ResNet50의 mAP(오른쪽). 출처.
다음으로, 검색에 대해 논의해 보겠습니다. 아시다시피, 검색의 목표는 주어진 쿼리에 대해 유사한 컨텍스트 집합을 찾는 것입니다. 따라서 MRL로 학습된 ResNet50과 고정 크기 특징 표현 간의 검색 품질을 mean Average Precision(mAP)을 사용해 테스트했습니다. 위 시각화에 설명된 것처럼, MRL 모델은 모든 특징 표현 크기에서 고정 크기 대응 모델에 비해 mAP가 최대 3% 향상되었습니다.
MRL을 사용하면 여러 특징 표현 크기를 활용할 수 있으므로, MRL로 정확도를 유지하면서 검색 프로세스를 어떻게 가속화할 수 있는지 살펴보는 것도 특히 흥미롭습니다. 이를 적응형 검색이라는 시나리오에서 테스트할 수 있습니다.
앞서 언급했듯이, 검색 프로세스에는 후보 선정과 재순위화라는 두 가지 주요 단계가 있습니다. 적응형 검색 설정은 200개의 후보 집합을 후보 선정하는 데 16차원 특징 표현을 사용하고, 재순위화에는 2048차원 특징 표현을 사용하는 방식이었습니다. 이 접근법은 후보 선정과 재순위화 모두에 2048차원 특징 표현을 사용하는 또 다른 방법과 비교됩니다.
ImageNet-1K에서 MRL 모델을 사용하면 이론적으로 검색 프로세스를 최대 128배까지 가속할 수 있습니다. 이를 이론적이라고 부르는 이유는 실제 애플리케이션에서는 검색 작업에 순진한 최근접 이웃 알고리즘을 사용하는 경우가 드물기 때문입니다. 대신 FAISS, ANNOY, HNSW와 같은 근사 최근접 이웃(ANN) 알고리즘이 검색 속도를 향상하기 위해 일반적으로 구현됩니다. 이 경우, MRL 모델을 사용하는 설정은 동일한 하드웨어에서 HNSW 알고리즘을 사용하는 검색 프로세스와 비교해 14배의 속도 향상을 달성합니다.
ImageNet-4K 데이터셋에서도 유사한 결과가 관찰되었으며, 논문의 저자들은 후보 목록 선정을 위해 64차원 특징 표현을, 재순위화를 위해 2048차원 표현을 사용하는 MRL 모델을 사용했습니다. 결과는 이론적으로 32배의 속도 향상과 실제 환경에서 6배의 속도 향상을 나타냅니다.
ImageNet-1K(왼쪽) 및 ImageNet-4K(오른쪽)에서 Adaptive Retrieval (AR)에 대한 mAP@10과 MFLOPs/Query 간의 트레이드오프.
ImageNet-1K(왼쪽) 및 ImageNet-4K(오른쪽)에서 Adaptive Retrieval (AR)에 대한 mAP@10과 MFLOPs/Query 간의 트레이드오프. 출처.
MRL 접근 방식 구현
이 섹션에서는 MRL 접근 방식의 간단한 구현을 살펴보겠습니다. 특히 NLI 데이터셋에서 MRL 접근 방식으로 학습된 MPNet base model을 사용하여 다양한 크기의 임베딩을 생성하겠습니다. 이미 알고 계실 수 있듯이, 원래의 MPNet base model은 768차원 임베딩만 생성할 수 있습니다. 그러나 이 학습된 MPNet 모델은 768, 512, 256, 128, 64차원 임베딩을 생성할 수 있으며, 학습 중 각 차원의 가중치는 1로 설정됩니다.
다음 코드를 통해 SentenceTransformers 라이브러리의 도움을 받아 이 모델을 사용할 수 있습니다.
from sentence_transformers import SentenceTransformer
matryoshka_dim_short = 64
matryoshka_dim_long = 768
text = ["The weather is so nice!"]
short_embedding = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim_short).encode(text)
long_embedding = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim_long).encode(text)
print(f"Shape: {short_embedding.shape, long_embedding.shape}")
print(short_embedding[0][0:10])
print(long_embedding[0][0:10])
"""
Output:
Shape: ((1, 64), (1, 768))
[-0.33891088 0.01647538 -0.29915053 0.24952686 -0.04321517 -0.31616145
-0.12996909 -0.05221268 0.02296597 0.07074839]
[-0.33891088 0.01647538 -0.29915053 0.24952686 -0.04321517 -0.31616145
-0.12996909 -0.05221268 0.02296597 0.07074839]
"""
위 코드에서는 학습된 MPNet base model을 사용하여 64와 768이라는 두 가지 다른 크기의 임베딩을 생성합니다. 각 차원의 가중치 또는 스케일링 계수가 1로 설정되어 있으므로, 위의 두 임베딩의 처음 10개 요소에서 볼 수 있듯이 두 임베딩은 동일한 요소를 가집니다.
각 요소의 값이 동일하므로 두 임베딩 간의 유사도는 1이 됩니다.
from sentence_transformers.util import cos_sim
similarities = cos_sim(short_embedding[0], long_embedding[0][:matryoshka_dim_short])
print(similarities)
# tensor([[1.]])
MRL 접근 방식으로 자체 모델을 학습하고 싶다면, 이 방법의 official GitHub repo를 살펴볼 수 있습니다.
결론
MRL 접근 방식은 머신 러닝에서 비용-성능 트레이드오프의 균형을 맞추는 솔루션을 도입합니다. 어떤 머신 러닝 모델이든 다양한 크기의 특징 표현을 생성할 수 있게 함으로써, MRL은 사용 사례와 리소스에 따라 속도 또는 정확도 중 하나에 최적화할 수 있는 유연성을 제공합니다. 우리는 이러한 적응성이 분류 및 검색과 같은 애플리케이션에서 가치가 있음을 확인했으며, 더 작은 표현과 더 큰 표현 사이를 전환하면 성능을 크게 희생하지 않으면서 효율성을 크게 향상할 수 있습니다.
실험 결과는 MRL이 텍스트, 비전, 멀티모달 작업을 포함한 여러 도메인에서 전통적인 고정 크기 모델의 정확도와 맞먹거나 종종 능가한다는 것을 보여줍니다. 또한 이 접근 방식이 ResNet50, ViT, BERT와 같은 기존 아키텍처와 호환된다는 점은 다양한 모델 모달리티 전반에 적용 가능하다는 것을 강조합니다. 실제 시나리오에서 상당한 속도 향상과 향상된 검색 품질의 잠재력을 갖춘 MRL은 더 효율적이고 다재다능한 머신 러닝 솔루션을 위한 유망한 발전을 나타냅니다.
관련 리소스
계속 읽기
VidTok: Rethinking Video Processing with Compact Tokenization
VidTok tokenizes videos to reduce redundancy while preserving spatial and temporal details for efficient processing.
Milvus WebUI: A Visual Management Tool for Your Vector Database
Explore Milvus WebUI to monitor, manage, and optimize your vector database with real-time insights, performance tracking, and system health monitoring.
DeepSeek Always Busy? Deploy It Locally with Milvus in Just 10 Minutes—No More Waiting!
Learn how to set up DeepSeek-R1 on your local machine using Ollama, AnythingLLM, and Milvus in just 10 minutes. Bypass busy servers and enhance AI responses with custom data.