Sentence Transformers 6.0 MultiVectorEncoder: 훈련 및 미세조정 가이드

요약

Hugging Face는 Sentence Transformers v6.0에 MultiVectorEncoder 모델 유형을 추가하고, 단일 RTX 3090에서 14.5시간 내에 ColBERT 스타일의 후기 상호작용 검색기(예: multi-vector-encoder/mLateOn-medical)를 미세조정하는 단계별 훈련 스크립트를 공개했습니다. 이 모델은 의료 검색 벤치마크에서 50개 이상의 밀도형, 희소형, 어휘형 및 다중 벡터 기반 모델 중에서 0.9139 NDCG@10의 최고 성능을 기록했습니다.

다중 벡터(후기 상호작용) 모델이란?

다중 벡터 모델은 전체 텍스트를 단일 벡터로 압축하는 대신, 각 토큰마다 하나의 임베딩을 유지합니다.

  • 검색은 MaxSim 연산자를 사용합니다: 각 쿼리 토큰이 가장 잘 매칭되는 문서 토큰을 찾고, 점수를 합산합니다.
  • 토큰 수준의 매칭은 단일 벡터 모델이 평균화하는 세부적인 관련성 신호를 유지하며, 일반적으로 더 강력한 검색 성능을 제공하지만 인덱스 크기가 커집니다.
  • 자세한 아키텍처, 인코딩, 점수 계산 및 인덱싱에 대한 내용은 "Sentence Transformers와 함께하는 다중 벡터(후기 상호작용) 임베딩 모델"이라는 보충 포스트를 참조하세요.

"밀도형 임베딩 모델은 전체 텍스트를 단일 벡터로 압축하고, 유사도는 두 개의 요약 간의 단일 내적 연산으로 계산됩니다. 다중 벡터 모델은 … 각 토큰마다 작은 벡터를 유지하고, MaxSim 연산자를 사용해 쿼리를 문서와 비교합니다." – 블로그 발췌

왜 다중 벡터 모델을 미세조정해야 할까?

미세조정은 모델을 특정 도메인(예: 의료, 법률, 코드)의 어휘, 쿼리 스타일, 관련성 개념에 맞게 조정합니다. 주요 관찰 사항:

  • 도메인 신호는 토큰 단위로 포착되므로, 적은 양의 도메인 내 데이터만으로도 큰 성능 향상이 가능합니다.
  • 대부분의 공개 체크포인트는 문서를 180~512 토큰으로 자릅니다. 긴 문장(의료 데이터셋 MIRIAD의 평균 941 토큰)에서는 이로 인해 최대 0.24 NDCG@10의 성능 손실이 발생할 수 있습니다.
  • 비지도 체크포인트(사전 대조적이지만 아직 지도 학습되지 않은)에서 시작하면 도메인 적응에 있어 완전히 지도된 체크포인트보다 일관되게 더 높은 성능을 보입니다.
  • 전체 미세조정 파이프라인은 단일 소비자 GPU에서 실행 가능하므로, 대부분의 팀이 도메인 특화 검색기를 접근할 수 있습니다.

훈련 구성 요소 개요

구성 요소 역할
모델 MultiVectorEncoder 인스턴스 – 사전 훈련된 체크포인트 또는 기본 트랜스포머 기반의 새 모델
데이터셋 datasets.Dataset 또는 DatasetDict – 쿼리-문서 쌍(또는 손실에 필요한 다른 형식) 포함
손실 배치 내 음성 샘플을 사용한 대조 손실(CachedMultiVectorMultipleNegativesRankingLoss) 또는 지식 증류 손실
훈련 인수 MultiVectorEncoderTrainingArguments – 배치 크기, 학습률, 프롬프트 등
평가기 NDCG@10, acc@1 등에 대한 MultiVectorInformationRetrievalEvaluator
트레이너 위 요소들을 조율하는 MultiVectorEncoderTrainer

아래 각 섹션은 독립적으로 읽을 수 있으며, 코드 조각은 pip install -U "sentence-transformers[train]"로 완전히 실행 가능합니다.

모델 선택 및 구성

기존 체크포인트 미세조정

from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192},  # 전체 길이 문서 허용
)
# 내장 길이 제한 제거 (예: 180~512 토큰)
model[0].query_length = None
model[0].document_length = None
# 선택 사항: 인덱스 크기 축소를 위해 구두점 토큰 건너뛰기
import string
model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)
  • 체크포인트는 이미 쿼리/문서 마커 토큰, 프로젝션 헤드, 점수 계산용 스킵리스트를 포함하고 있습니다.
  • 제한을 제거하면 모델이 최대 1,400 토큰의 의료 문장 전체를 처리할 수 있습니다.

기본 트랜스포머에서 새 모델 구축

from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder(
    "answerdotai/ModernBERT-base",
    model_kwargs={"torch_dtype": "float32"},
)
  • 파이프라인은 무작위 토큰 수준 프로젝션(128차원)과 일반적인 ColBERT 모듈을 추가합니다.
  • 강력한 밀도형 백본(예: Alibaba-NLP/gte-modernbert-base)도 25,000개의 쌍으로 훈련한 후 체크포인트 성능에 근접합니다.

어떤 시작점 선택할까?

25,000개의 MIRIAD 쌍을 기반으로 한 경험적 비교:

체크포인트 제로샷 NDCG@10 미세조정 후 Δ
lightonai/mLateOn-unsupervised 0.9087 0.9398 +0.0311
lightonai/mLateOn 0.9277 0.9319 +0.0042
lightonai/LateOn-unsupervised 0.9026 0.9206 +0.0180
lightonai/LateOn 0.9185 0.9105 –0.0080
lightonai/GTE-ModernColBERT-v1 0.9198 0.9007 –0.0191
gte-modernbert-base에 새 헤드 – 0.9177 –
결론: 비지도 훈련된 체크포인트가 가장 잘 적응하며, 완전히 지도된 체크포인트는 종종 성능 저하를 겪습니다.

데이터셋 준비

트레이너는 어떤 datasets.Dataset(Hub 또는 로컬)도 수용합니다. 필요 형식은 손실에 따라 달라집니다:

  • 손실이 감독 신호를 기대하는 경우 라벨 컬럼(label 또는 score)이 필요합니다.
  • 입력 컬럼은 순서가 있으며, 첫 번째 컬럼은 쿼리로 간주되며, 이후 컬럼은 문서로 간주됩니다(단, router_mapping로 재정의 가능).

Hub에서 로드하기 (MIRIAD 예시)

from datasets import load_dataset
train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train")
print(train_dataset)
# Dataset({features: ['question', 'passage_text'], num_rows: 4_467_542})
  • 각 행은 (질문, passage_text) 쌍을 제공하며, 문장의 평균 길이는 941 토큰입니다.

로컬 CSV/JSON 예시

from datasets import load_dataset
dataset = load_dataset("csv", data_files="my_file.csv")
# 또는 JSON
# dataset = load_dataset("json", data_files="my_file.json")
  • 사용자 정의 전처리가 필요한 경우 Dataset.from_dict를 사용하세요.

손실 함수 선택

질문-답변 쌍의 경우 추천 손실은 배치 내 음성 샘플과 GradCache를 사용하는 것입니다:

from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss
loss = CachedMultiVectorMultipleNegativesRankingLoss(
    model=model,
    mini_batch_size=16,  # 메모리용 청크 크기; 블로그에서는 효과적 배치 128 사용
)
  • mini_batch_size는 메모리 제어에 사용되며, 효과적인 대조 배치 크기는 여전히 128입니다.
  • 스케일 파라미터: MaxSim 점수에 대해 기본값 scale=1.0을 유지하세요. 밀도형 기본값(scale=20.0)을 사용하면 그래디언트가 포화됩니다.
  • 증류를 위한 경우 MultiVectorDistillKLDivLoss를 참조하세요(블로그 예시에서는 사용되지 않음).

훈련 인수

최고의 결과를 얻은 주요 설정:

from sentence_transformers import MultiVectorEncoderTrainingArguments
from sentence_transformers.base.sampler import BatchSamplers
args = MultiVectorEncoderTrainingArguments(
    output_dir="models/mLateOn-medical",
    num_train_epochs=1,
    per_device_train_batch_size=128,   # GradCache를 통한 효과적 배치
    per_device_eval_batch_size=16,
    learning_rate=1e-4,
    warmup_steps=0.05,
    prompts={"question": "[Q] ", "passage_text": "[D] "},
    fp16=False,
    bf16=True,
    batch_sampler=BatchSamplers.NO_DUPLICATES,
    eval_strategy="steps",
    eval_steps=0.1,
    save_strategy="steps",
    save_steps=0.05,
    logging_steps=0.01,
    run_name="mLateOn-medical",
)
  • 프롬프트는 명시적으로 제공되어야 하며, 모델은 저장된 마커를 자동으로 적용하지 않습니다.
  • max_length를 설정하지 않으면 훈련이 전체 문서 길이를 볼 수 있습니다.
  • 5e-6 → 2e-4 범위에서 탐색한 결과, 더 높은 학습률(1e-4)이 가장 우수한 성능을 보였습니다.

검색용 평가기

가장 유의미한 평가기는 보류된 쿼리 세트와 방해물이 포함된 코퍼스를 기반으로 만든 MultiVectorInformationRetrievalEvaluator입니다:

from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator
# 코퍼스, 쿼리, 관련 문서 매핑 구축 (전체 루프는 블로그 참조)
evaluator = MultiVectorInformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant_docs,
    name="miriad-dev",
    batch_size=16,
)
  • 포화를 피하기 위해 어려운 방해 문장을 포함하세요. 블로그에서는 10,000개의 골드 세트에 약 19만 개의 랜덤 훈련 문장을 추가했습니다.

전체 훈련 스크립트

다음 스크립트는 mLateOn-medical 모델을 재현합니다:

import logging, string, traceback
from datasets import load_dataset
from sentence_transformers import (
    MultiVectorEncoder,
    MultiVectorEncoderModelCardData,
    MultiVectorEncoderTrainer,
    MultiVectorEncoderTrainingArguments,
)
from sentence_transformers.base.sampler import BatchSamplers
from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator

logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO)

def main():
    # 1️⃣ 비지도 체크포인트 로드 및 제한 제거
    model = MultiVectorEncoder(
        "lightonai/mLateOn-unsupervised",
        model_kwargs={"torch_dtype": "float32"},
        processor_kwargs={"model_max_length": 8192},
        model_card_data=MultiVectorEncoderModelCardData(
            language="en",
            license="apache-2.0",
            model_name="mLateOn finetuned on MIRIAD medical retrieval",
        ),
    )
    model[0].query_length = None
    model[0].document_length = None
    model[2].skiplist_words = list(string.punctuation)
    model[2].resolve_with_tokenizer(model.tokenizer)

    # 2️⃣ 100만 개의 의료 QA 쌍 로드
    train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train").select(range(1_000_000))

    # 3️⃣ 손실 정의 (GradCache 배치 내 음성 샘플)
    loss = CachedMultiVectorMultipleNegativesRankingLoss(model=model, mini_batch_size=16)

    # 4️⃣ 가벼운 개발 평가기 (500개 쿼리)
    eval_split = load_dataset("tomaarsen/miriad-4.4M-split", split="eval")
    corpus, queries, relevant_docs, passage_to_id = {}, {}, {}, {}
    for idx, row in enumerate(eval_split):
        if row["passage_text"] not in passage_to_id:
            pid = f"p{len(passage_to_id)}"
            passage_to_id[row["passage_text"]] = pid
            corpus[pid] = row["passage_text"]
        if idx < 500:
            qid = f"q{idx}"
            queries[qid] = row["question"]
            relevant_docs[qid] = {passage_to_id[row["passage_text"]]}  
    dev_evaluator = MultiVectorInformationRetrievalEvaluator(
        queries=queries, corpus=corpus, relevant_docs=relevant_docs, name="miriad-dev", batch_size=16
    )

    # 5️⃣ 훈련 인수 (이전 섹션 참조)
    args = MultiVectorEncoderTrainingArguments(
        output_dir="models/mLateOn-medical",
        num_train_epochs=1,
        per_device_train_batch_size=128,
        per_device_eval_batch_size=16,
        learning_rate=1e-4,
        warmup_steps=0.05,
        prompts={"question": "[Q] ", "passage_text": "[D] "},
        fp16=False,
        bf16=True,
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.05,
        logging_steps=0.01,
        run_name="mLateOn-medical",
    )

    # 6️⃣ 트레이너 및 훈련
    trainer = MultiVectorEncoderTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        loss=loss,
        evaluator=dev_evaluator,
    )
    trainer.train()

    # 7️⃣ 저장 및 선택적으로 Hub에 업로드
    model.save_pretrained("models/mLateOn-medical/final")
    try:
        model.push_to_hub("mLateOn-medical")
    except Exception:
        logging.error("모델 업로드 실패: \n" + traceback.format_exc())

if __name__ == "__main__":
    main()
  • 실행 시간: 단일 RTX 3090에서 14.5시간 (피크 17.5GB VRAM).
  • 데이터 효율성: 10만 개의 쌍(약 75분)만으로도 전체 백만 쌍 훈련 결과와 NDCG@10 차이가 0.012 이내입니다.

인덱스 크기 및 최적화

다중 벡터 인덱스는 각 토큰이 벡터를 생성하기 때문에 더 큽니다(941토큰 문장당 약 878개 벡터). 20만 개 문장의 원시 fp16 저장 용량은 약 45GB입니다.

토큰 풀링

HierarchicalTokenPooling(pool_factor=4)는 벡터 수를 4배 줄이며, NDCG@10 손실은 < 0.0033입니다.

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
pooling = HierarchicalTokenPooling(pool_factor=4)
embeddings = model.encode_document(passages, token_pooling=pooling)
  • 벡터 수의 ¼ → 약 11GB, NDCG@10 ≈ 0.8991.

양자화 및 절단 (PLAID)

1비트 잔차 양자화와 적절한 절단을 사용하면:

구성 유지된 벡터 인덱스 크기 NDCG@10
1비트 PLAID, 모든 벡터 100% 3.37GB 0.8984
1비트 PLAID + 절단 65% 2.23GB 0.8830
1비트 PLAID + 절단 42% 1.45GB 0.8642
  • 양자화는 13배 이상의 용량 감소를 제공하며, NDCG 손실은 0.02 이내로, 다중 벡터 검색이 밀도형 모델과 유사한 저장 용량을 갖게 합니다.

평가 결과

미세조정된 multi-vector-encoder/mLateOn-medical vs. 50개 이상의 기준 모델, 20만 문장 의료 벤치마크(1,000개 보류된 쿼리):

모델 가족 NDCG@10 acc@1
multi-vector-encoder/mLateOn-medical 다중 벡터(미세조정) 0.9139 0.849
lightonai/mLateOn 다중 벡터(제로샷) 0.8520 0.758
lightonai/GTE-ModernColBERT-v1 (제한 제거) 다중 벡터(제로샷) 0.8502 0.763
Qwen/Qwen3-Embedding-4B 밀도형(제로샷) 0.7817 0.669
voyageai/voyage-4-nano 밀도형(제로샷) 0.7563 0.638
BM25 어휘형 0.7501 0.641
naver/splade-v3 희소형(제로샷) 0.6853 0.574
결론: *문서 길이가 길 때 후기 상호작용 모델이 우세하며, 미세조정은 가장 강력한 제로샷 모델보다 +0.062 NDCG@10의 추가 성능 향상을 제공합니다.*

실용적 통찰

  1. 비지도 체크포인트(예: lightonai/mLateOn-unsupervised)에서 시작하여 도메인 미세조정을 수행하세요.
  2. 문서 길이 제한을 제거하여 데이터 길이에 맞추세요. 그렇지 않으면 긴 문장에서 최대 0.24 NDCG@10의 성능 손실이 발생합니다.
  3. **GradCache(CachedMultiVectorMultipleNegativesRankingLoss)**를 사용하여 단일 GPU에서 큰 효과적 배치 크기를 달성하세요.
  4. 구두점 스킵리스트를 적용하여 인덱스 크기를 약 10% 줄이되 품질 손실 없이 진행하세요.
  5. 인덱스 압축(토큰 풀링, PLAID 양자화)에 투자하여 다중 벡터 저장 용량을 밀도형 모델과 유사한 수준으로 유지하세요.
  6. 적은 데이터(10만 쌍)만으로도 강력한 성능을 얻을 수 있으므로, 많은 조직이 이 접근법을 활용할 수 있습니다.

추가 자료

  • 훈련 예제: MIRIAD 의료 미세조정, MS MARCO 대조 및 증류, 다중 모달 ColPali, PEFT LoRA 어댑터.
  • 문서화: 설치, 빠른 시작, 사용법, 사용자 정의 모델 생성, 사전 훈련된 모델 목록, 훈련 및 손실 개요, API 참조.
  • 보충 포스트(사용법): Sentence Transformers와 함께하는 다중 벡터(후기 상호작용) 임베딩 모델 – 인코딩, 인덱싱, 배포에 대해 다룹니다.

"다중 벡터 인덱스가 너무 크다는 주장은 적절하게 구성된 인덱스에서는 성립하지 않습니다." – 오마르 카타브에게 양자화 측정에 감사드립니다.

Sources

관련