Sentence Transformers를 사용한 재정렬 모델 훈련 및 파인튜닝
Hugging Face는 Sentence Transformers 라이브러리를 사용하여 재정렬 모델(Cross Encoder)을 훈련하고 파인튜닝하기 위한 프레임워크를 자세히 설명하고 있습니다. 도메인 특화 데이터로 훈련함으로써 개발자는 작고 효율적인 재정렬기를 만들 수 있으며, 이는 특수 검색 작업에서 훨씬 큰 범용 모델보다 뛰어난 성능을 보입니다.
재정렬 모델 vs. 임베딩 모델
Cross Encoder 재정렬 모델은 쿼리와 문서와 같은 텍스트 쌍의 관련성을 공유 신경망을 통해 함께 처리하여 단일 출력 점수를 생성함으로써 평가합니다. 이는 텍스트를 독립적으로 벡터에 임베딩하고 거리 메트릭을 통해 유사성을 계산하는 bi-encoder(임베딩 모델)와 다릅니다.
Cross Encoder는 모든 가능한 쌍을 처리해야 하므로 계산 비용이 더 많이 들지만 일반적으로 더 높은 정확도를 제공합니다. 따라서 고성능 검색을 위한 산업 표준은 두 단계의 "retrieve and rerank" 파이프라인이며, 초기 검색을 위해 빠른 임베딩 모델을 사용하고 상위 k 결과를 정제하기 위해 Cross Encoder를 사용합니다.
핵심 훈련 구성 요소
재정렬 모델을 훈련하는 데는 CrossEncoderTrainer를 통해 통합되는 다섯 가지 주요 구성 요소가 포함됩니다:
1. 데이터셋
CrossEncoderTrainer는 datasets.Dataset 또는 datasets.DatasetDict 객체를 지원합니다. 데이터는 Hugging Face Hub 또는 로컬 파일(CSV, JSON, Parquet, Arrow, SQL)에서 가져올 수 있습니다. 선택한 손실 함수와 호환성을 보장하려면 데이터셋에 레이블용으로 "label", "labels", "score", "scores" 중 하나의 이름을 가진 열이 있어야 하며, 나머지 열은 손실 함수가 요구하는 입력 수와 일치해야 합니다.
2. 하드 네거티브 마이닝
모델 성능은 네거티브의 품질에 크게 좌우됩니다. "soft negatives"는 완전히 무관한 경우이며, "hard negatives"는 관련 있어 보이지만 실제로는 아닌 구절입니다. Sentence Transformers는 이러한 어려운 예시를 식별하기 위해 mine_hard_negatives 함수를 제공하며, 이를 통해 모델이 구분을 더 정확히 하도록 강제합니다.
3. 손실 함수
손실 함수는 사용 가능한 데이터를 기반으로 최적화 과정을 안내합니다. LambdaLoss나 ListNetLoss와 같은 학습-순위 손실이 존재하지만, BinaryCrossEntropyLoss는 라벨이 있는 쌍에 대해 매우 효과적이고 간단한 옵션으로 남아 있습니다.
4. 훈련 인자
맞춤 설정은 CrossEncoderTrainingArguments를 통해 처리되며, 개발자는 학습률, 배치 크기, 워밍업 비율 및 정밀도 설정(FP16/BF16)을 조정할 수 있습니다. batch_sampler=BatchSamplers.NO_DUPLICATES 설정은 배치 내 네거티브를 사용하는 손실에 특히 권장됩니다.
5. 평가자
단순 손실을 넘어 성능을 추적하기 위해 Sentence Transformers는 여러 내장 평가자를 제공합니다:
CrossEncoderClassificationEvaluator: 이진 또는 다중 클래스 레이블용.CrossEncoderCorrelationEvaluator: 유사도 점수용(예: STSb 데이터셋 사용).CrossEncoderRerankingEvaluator: 쿼리, 양성, 음성을 사용한 재정렬 성능 평가용.CrossEncoderNanoBEIREvaluator: 영어 재정렬을 위한 경량 평가자.
다중 데이터셋 훈련
CrossEncoderTrainer는 서로 다른 형식이거나 서로 다른 손실 함수를 필요로 하는 경우에도 여러 데이터셋을 동시에 훈련할 수 있습니다. 이는 데이터셋 사전과 선택적인 손실 함수 사전을 통해 관리됩니다. 샘플링 전략은 다음과 같습니다:
ROUND_ROBIN: 각 데이터셋에서 동일하게 샘플링하며 하나가 소진될 때까지.PROPORTIONAL: 각 데이터셋 크기에 비례하여 샘플링하여 모든 샘플이 사용되도록 보장.
성능 평가 및 결과
실제 적용 사례에서, ModernBERT-base를 기반으로 한 재정렬기는 GooAQ 데이터셋의 99k 쿼리-답변 쌍에 대해 파인튜닝되었습니다. BinaryCrossEntropyLoss와 하드 네거티브 마이닝을 사용한 결과 모델(tomaarsen/reranker-ModernBERT-base-gooaq-bce)은 13개의 일반적인 오픈소스 재정렬기보다 뛰어났으며, 크기가 최대 4배인 모델도 포함되었습니다.
GooAQ NDCG@10 결과 (상위 30 재정렬)
| 모델 | 파라미터 | 실제 NDCG@10 | 평가 NDCG@10 |
|---|---|---|---|
| Retriever only (No reranking) | - | 59.12 | 59.12 |
| BAAI/bge-reranker-large | 560M | 73.20 | 77.46 |
| mixedbread-ai/mxbai-rerank-large-v2 | 1.54B | 75.40 | 80.04 |
| ModernBERT-base-gooaq-bce | 150M | 77.14 | 83.51 |
| ModernBERT-large-gooaq-bce | 396M | 79.42 | 85.81 |
기술 훈련 팁
- 과적합 방지: Cross Encoder는 빠르게 과적합됩니다.
load_best_model_at_end와metric_for_best_model을 사용한 평가자를 이용해 최고 성능 모델을 저장하세요. - 네거티브 균형 맞추기: 하드 네거티브만 사용하면 쉬운 작업에서 성능이 저하될 수 있습니다. 무작위 네거티브와 하드 네거티브를 혼합하면 이 문제를 완화할 수 있습니다.
- 효율성: 도메인 특화 데이터에 작은 재정렬기를 파인튜닝하면 대규모 범용 모델을 사용하는 것보다 검색 정확도를 향상시키고 추론 지연 시간을 줄일 수 있습니다.