10억 개의 학습 쌍을 사용한 Hugging Face 문장 임베딩 모델

TL;DR

Hugging Face는 최대 10억 개의 문장 쌍으로 구성된 대규모 코퍼스를 학습하여 최첨단 범용 문장 임베딩 모델을 개발했습니다. 이는 JAX/Flax 프레임워크와 Google TPU v3-8 인프라를 사용하여 대규모 대조 학습(contrastive learning)을 최적화함으로써 달성되었습니다.

학습 방법론

모델 아키텍처

문장 임베딩 모델은 클러스터링, 텍스트 마이닝, 질의응답과 같은 애플리케이션을 위해 문장의 의미론적 의미를 포착하기 위해 문장을 실수 벡터로 매핑합니다. 가능한 모든 문장의 집합은 무한하기 때문에, 이러한 모델은 구성 모듈(composition module)—일반적으로 문맥화된 단어 벡터에 대한 풀링 연산이 뒤따르는 Transformer—를 사용하여 최종 표현을 계산합니다.

Multiple Negative Ranking Loss (MNRL)

모델은 Multiple Negative Ranking Loss(InfoNCE 또는 NTXentLoss라고도 함)로 알려진 대조 학습 방법을 사용하여 학습되었습니다. 이 접근 방식은 인배치 네거티브(in-batch negatives)를 사용하여 임베딩 공간을 최적화합니다:

데이터셋 구성: 학습 세트는 의미가 유사한 쌍 $(a_i, p_i)$ (예: 질의-답변 쌍, 중복 질문, 또는 인용된 논문 제목)로 구성됩니다.
목표: 모델은 양성 쌍 $(a_i, p_i)$를 가까운 벡터로 매핑하는 동시에, $i \neq j$인 일치하지 않는 쌍 $(a_i, p_j)$를 먼 벡터로 밀어내는 방식으로 학습됩니다.
유사도 함수: 모델은 배치 내의 모든 쌍 사이의 유사도 행렬을 계산합니다. 사용되는 유사도 함수는 Cosine-Similarity 또는 Dot-Product입니다.
- Cosine-Similarity: 정규화된 벡터는 가장 높은 유사도가 자기 자신과의 유사도(1)임을 보장하며, 유클리드 거리와 비례하므로 k-means 클러스터링과 호환됩니다.
- Dot-Product: 일부 근사 최근접 이웃(approximate nearest neighbor) 방법과 함께 사용하면 더 빠를 수 있지만, k-means 클러스터링과 작동하지 않으며 다른 벡터가 자기 자신과의 유사도보다 더 높은 내적값을 가질 수 있게 합니다.

점수 차이가 너무 작아지는 것을 방지하기 위해, 유사도 점수에 스케일링 인자 $C$ (일반적으로 $C=20$)가 적용됩니다: $sim_{scaled}(a, b) = C * sim(a, b)$.

임베딩 품질 최적화

배치 크기와 Hard Negatives

배치 구성은 대조 학습 성능에 매우 중요합니다. Hugging Face는 품질을 향상시키기 위한 세 가지 주요 레버를 확인했습니다:

Batch Size: 배치 크기가 클수록 일반적으로 모델 성능이 향상됩니다.
Hard Negatives: 양성 샘플 $p_i$와 의미론적으로 유사하지만 올바른 매칭이 아닌 샘플 $p_j$ (예: "What is the capital of France?" vs. "What is the capital of the US?")를 포함하면 모델이 더 정밀한 의미론적 구분을 학습하도록 강제합니다.
Cross-Dataset Batches: 단일 배치 내에 최소 두 개의 서로 다른 데이터셋을 혼합함으로써, 모델은 단일 주제 내의 국소적 구조뿐만 아니라 서로 다른 주제 간의 전역적 구조를 학습합니다.

인프라 및 데이터 규모

이 프로젝트는 대규모 대조 학습에 필요한 행렬 곱셈을 처리하기 위해 7개의 TPU v3-8을 활용했습니다. 학습 데이터는 총 10억 개의 문장 쌍에 달하는 연결된 데이터셋으로 구성되었습니다.

결과 및 애플리케이션

Hugging Face는 Mini-LM, RoBERTa, DistilBERT, MPNet을 포함한 아키텍처를 기반으로 20개의 범용 Sentence Transformer 모델을 학습시켜, 여러 범용 문장 유사도 평가 작업에서 최첨단(SOTA) 결과를 달성했습니다.

모델과 함께 Sentence-Similarity, Question Answering, Gender Evaluation을 위한 8개의 특화된 데이터셋이 공개되었습니다. 이러한 임베딩은 다음과 같은 여러 실용적인 애플리케이션을 가능하게 합니다:

Sentence Similarity: 코사인 유사도를 사용하여 두 텍스트의 의미론적 유사성을 비교합니다.
Asymmetric QA: 후보 패시지가 특정 질의에 대한 답변인지 여부를 판단합니다.
Search and Clustering: 내적(dot-product) 거리를 사용하여 질의에 대한 근처의 답변을 검색합니다.
Gender Bias Evaluation: 직업 기반 앵커 텍스트에서 성별이 드러나는 대명사를 위한 모델 유사도 점수를 비교함으로써 학습 세트에 내재된 성별 편향을 식별합니다.