10억 개 데이터 기반 문장 임베딩 학습법

문장을 의미를 담은 벡터로 매핑하는 문장 임베딩(Sentence Embedding) 기술을 고도화하기 위해, 10억 개의 학습 쌍을 활용한 대규모 학습 프로젝트 결과와 방법론을 공개한다.

주요 기술적 특징:

• 인프라 및 프레임워크: JAX/Flax 프레임워크와 7개의 TPU v3-8을 사용하여 효율적인 대규모 학습을 수행했다.
• 모델 구조: Transformer 모델을 기반으로 하며, 문맥화된 단어 벡터에 풀링(Pooling) 연산을 적용하여 최종 문장 표현을 도출한다.
• 학습 전략: Multiple Negative Ranking Loss(InfoNCE)를 사용하는 대조 학습(Contrastive Learning) 방식을 채택했다. 이는 배치 내의 정답 쌍(Positive pairs)은 유사도를 높이고, 나머지 샘플(In-batch negatives)과는 거리를 벌리는 방식이다.
• 유사도 함수: 코사인 유사도와 내적(Dot-product)의 특성을 비교하며, 실제 학습 시에는 점수 차이를 조절하기 위해 스케일링 인자를 적용했다.