Hugging Face 팀이 10억 개의 문장 쌍으로 Sentence Transformer 모델을 학습하면서 대조 학습(Contrastive Learning)과 배치 최적화를 통해 여러 일반 목적 문장 임베딩 모델에서 SOTA(State-of-the-Art) 달성

Context

문장을 의미론적으로 포착하는 벡터로 매핑하는 문장 임베딩 모델은 클러스터링, 텍스트 마이닝, 질의응답 등 다양한 하위 작업에 사용되어야 하지만, 기존 모델들의 성능 수준이 최고 수준에 도달하지 못했다.

Technical Solution

• Multiple Negative Ranking Loss 채택: 배치 내 모든 문장 쌍에 대해 유사도 행렬을 계산한 후 대조 학습으로 의미론적으로 관련된 쌍(a_i, p_i)은 높은 유사도, 무관한 쌍(a_i, p_j, i≠j)은 낮은 유사도를 할당하는 손실 함수 적용
• 유사도 함수 선택 및 스케일링: Cosine-Similarity 또는 Dot-Product 중 선택하여 사용하되, 점수 차이가 작은 문제를 해결하기 위해 스케일 계수 C=20을 곱해 유사도 스케일링 적용
• 배치 크기 확대: 대조 학습 성능 향상을 위해 더 큰 배치 크기 사용
• Hard Negatives 포함: 배치에 의미론적으로 유사하지만 구분이 어려운 음성 샘플(hard negatives)을 의도적으로 포함하여 모델 식별 능력 강화
• JAX/Flax 프레임워크와 7개의 TPU v3-8 인프라 활용: 효율적인 딥러닝 학습을 위해 Google의 Flax/JAX 프레임워크 채택 및 고성능 하드웨어 인프라 구성
• 20개의 일반 목적 Sentence Transformer 모델 학습: Mini-LM, RoBERTa, DistilBERT, MPNet 등 다양한 기반 모델 아키텍처로 모델 생성

Impact

학습된 모델들이 여러 일반 목적 문장 유사도 평가 작업(evaluation tasks)에서 SOTA 달성.

Key Takeaway

대조 학습에서 배치 구성(크기 + hard negatives 포함)이 모델 성능에 핵심적인 영향을 미치므로, 규모 있는 임베딩 모델 개발 시 단순히 데이터량 증가뿐 아니라 배치 내 샘플 선택 전략이 중요한 설계 요소다.