Sentence Transformers를 사용한 다중 벡터 임베딩 모델 학습 및 미세 조정

MultiVectorEncoder에 대한 내용과 ColBERT 스타일의 후기 상호작용(late interaction) 검색을 위한 완전한 학습 접근 방식을 소개합니다. 이 블로그 게시물에서는 이를 사용하여 일반 목적 검색기보다 귀하의 데이터에서 성능이 우수한 다중 벡터 모델을 미세 조정하는 방법을 보여드리겠습니다. 또한 이 방법은 처음부터 강력한 새로운 다중 벡터 모델을 학습할 수도 있습니다. 아래 내용은 `pip install -U

밀집 임베딩 모델(dense embedding model)은 전체 텍스트를 단일 벡터로 압축하며, 유사도는 두 요약 간의 내적(dot product) 값 중 하나입니다. 반면, 다중 벡터 모델(multi-vector model)(후기 상호작용 또는 ColBERT 스타일 모델이라고도 함)은 이러한 압축 과정을 건너뜁니다. 이 모델은 토큰당 하나의 작은 벡터를 유지하며, 쿼리 토큰이 가장 잘 맞는 문서 토큰을 찾아내고 그 점수들을 합산하는 MaxSim 연산자를 사용하여 쿼리와 문서를 점수화합니다. 토큰 수준의 매칭은 단일 벡터가 평균화해야 하는 미세한 신호(fine-grained signals)를 정확하게 보존하므로, 일반적으로 더 강력한 검색 성능을 제공하지만 인덱스 크기가 커진다는 단점이 있습니다.

동반되는 Multi-Vector Embedding Models 블로그 게시물에서 아키텍처, 인코딩, 점수화 및 인덱싱에 대해 자세히 다루었으므로, 여기서는 간략하게 설명하고 학습 과정으로 넘어가겠습니다.

다중 벡터 모델을 미세 조정(Finetuning)하면 특정 도메인에서의 검색 성능이 크게 향상됩니다. 웹 검색, 법률 문서 조사, 코드 검색, 과학 문헌 검토 등 각 분야마다 어휘, 쿼리 스타일, 관련성 개념이 모두 다르기 때문입니다. 쿼리와 문서를 토큰별로 매칭하기 때문에, 다중 벡터 모델은 단일 벡터 모델이 평균화하는 경향이 있는 미세한 도메인 신호를 포착하며, 적지 않은 양의 인-도메인(in-domain) 미세 조정 데이터만으로도 매우 좋은 성능을 보입니다.

그 외에도, 대부분 출시된 검색 모델들은 짧은 구절(short passages)에 맞춰 설정되었습니다. 고전적인 ColBERT 체크포인트는 문서를 180개 또는 300개의 토큰에서 잘라내고, 많은 인기 있는 밀집 모델은 MS MARCO 스타일의 학습 데이터가 그 이상으로 거의 가지 않기 때문에 256개 또는 512개로 자릅니다. 만약 문서가 길다면, 이 모델들은 점수를 매기기 전에 모든 문서의 대부분을 조용히 버립니다. 평균 941개의 토큰인 구절로 진행한 저의 의학 평가에서, 저는 이러한 잘라냄(truncation)이 최대 0.24 NDCG@10만큼 손실시킨다는 것을 측정했는데, 이는 모델 아키텍처 간의 어떤 차이보다도 훨씬 큰 수치입니다. 자체 모델을 학습시킬 때는 사용자 데이터가 필요로 하는 문서 길이를 설정해야 합니다.

LightOn도 코드 검색에서 이와 같은 역동성을 겪었으며, 일반적인 LateOn만으로는 충분하지 않아 LateOn-Code를 학습시켰습니다. 귀하의 도메인(의료, 법률, 금융 또는 회사 내부 문서 등)에 맞는 공식 모델은 없습니다. 본 블로그 게시물에서는 단일 소비자 GPU에서 몇 시간 만에 직접 구축하는 방법을 보여드립니다.

MultiVectorEncoder 모델을 학습시키려면 다음 구성 요소들이 필요합니다:

모델(Model): 미세 조정할 모델 또는 새로 구축할 아키텍처입니다.데이터셋(Dataset): 학습 및 평가에 사용되는 데이터입니다.손실 함수(Loss Function): 모델의 성능을 측정하고 최적화 과정을 안내하는 함수입니다.학습 인자(Training Arguments) (선택 사항): 학습 성능, 추적 및 디버깅에 영향을 미치는 매개변수입니다.평가기(Evaluator) (선택 사항): 학습 전, 중 또는 후에 모델을 평가하기 위한 클래스입니다.트레이너(Trainer): 모든 학습 구성 요소를 하나로 통합합니다.

각 구성 요소들을 자세히 살펴보겠습니다.

Multi-vector 학습은 실제 시작 지점을 선택할 수 있는 자유를 제공하며, 이는 예상보다 더 중요합니다.

기존의 Multi-vector 모델을 추가적으로 미세 조정하고 싶다면, 아키텍처에 대해서는 전혀 걱정할 필요가 없습니다:

from sentence_transformers import MultiVectorEncoder
# 메모리가 처리할 수 있다면 fp32로 로드하는 것이 학습에 선호됩니다
model = MultiVectorEncoder(
...

체크포인트는 자체 레시피를 가지고 옵니다: 쿼리 및 문서 마커 토큰, 투영 헤드(projection head), 점수 스키플리스트(scoring skiplist) 등이 있습니다. 미세 조정을 위해서는 일반적으로 이 모든 것을 유지하고 데이터가 요구하는 부분만 변경하면 됩니다. 가장 먼저 확인해야 할 것은 길이 설정입니다. 많은 공개된 체크포인트가 문서를 180개에서 512개의 토큰으로 제한하기 때문에(Why Finetune? 참조), 제 의료 자료는 최대 1,400개의 토큰에 달합니다. mLateOn 계열은 이미 백본의 전체 8192 토큰 컨텍스트를 제공하지만, 시작 체크포인트가 제한을 가지고 있다면 이를 해제하십시오:

모델이 학습된 길이의 캡(caps) 대신 전체 문서를 읽도록 하십시오.
예를 들어, GTE-ModernColBERT-v1은 query_length=48 및 document_length=300으로 제공됩니다.
model[0].query_length = None
...

per-task caps가 설정되지 않으면, 잘림(truncation)은 토크나이저의 model_max_length로 폴백합니다. 이것이 제가 위에서 로드 시 이 제한을 구성한 이유입니다.

저는 또한 구두점 토큰을 문서 측 점수 계산 및 저장에서 제외하는 구두점 skiplist를 추가했습니다. 4가지 비교 실험(없음, 구두점, 불용어, 둘 다)에서 품질 면에서 약간 우위를 보였으며, 이 데이터셋에서 문서 인덱스를 무료로 9.6% 축소합니다:

import string
# model[2]는 MultiVectorMask 모듈입니다.
model[2].skiplist_words = list(string.punctuation)
...

또한 MultiVectorEncoder를 아무 기본 트랜스포머에도 연결할 수 있으며, 새로운 무작위 초기화된 토큰 레벨 투영이 추가됩니다:

from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder("answerdotai/ModernBERT-base", model_kwargs={"torch_dtype": "float32"})
# MultiVectorEncoder(
...

이것이 전형적인 ColBERT 파이프라인입니다. 문맥화된 토큰 임베딩을 생성하는 Transformer, 각 토큰을 128차원으로 투영하는 토큰 레벨의 Dense, 점수 계산 중 어떤 토큰을 포함할지 결정하는 MultiVectorMask, 그리고 토큰 레벨의 Normalize입니다. 투영은 무작위로 시작하므로, 이 모델이 유용해지기 전에 훈련이 필요합니다. 흥미롭게도, 이것은 강력한 밀집 임베딩 백본과도 작동합니다. Alibaba-NLP/gte-modernbert-base에 대한 새로운 투영은 제가 실험한 기존 체크포인트 시작점들과 0.03 이내의 값을 달성했으며, 오직 투영과 25k 개의 훈련 쌍만 사용했습니다.

전형적인 ColBERT 토크나이징 트릭([MASK] 질의 확장, [Q] / [D] 접두사 토큰, 문서 길이 제한, 구두점 skiplist)은 기본적으로 비활성화되어 있으며 구성 가능합니다. 전체 목록은 Creating Custom Models를 참조하십시오. 참고로, 저는 [MASK]

제 도메인에 맞춰 파인튜닝하기 위해 네 가지 구성으로 쿼리 확장을 시도했지만, 어느 것도 측정 가능한 차이를 만들지 못했습니다. 그러니 고전적인 레시피를 사용해야 한다는 부담을 느끼실 필요는 없습니다.

저는 이 블로그 게시물을 준비하면서 직접 측정했는데, 여섯 개의 시작 지점을 가져와 MIRIAD의 25k 의료 질문-문단 쌍으로 동일한 레시피로 각각 학습시킨 후, 1,000개의 홀드아웃(held-out) 질문을 50,000개 문단 코퍼스에 대해 평가했습니다:

이 결과는 저를 놀라게 했고, 두 모델 계열에서 재현되었습니다. *-unsupervised 체크포인트가 완성된 형제들보다 새로운 도메인에 훨씬 잘 적응하며, 시작점은 낮음에도 불구하고 그들을 추월했습니다. 이 체크포인트들은 대규모 대비 사전 학습(contrastive pretraining)을 거쳤지만 일반 검색(general retrieval)에 대한 지도 학습 파인튜닝 전의 상태이므로, 도메인 훈련 과정에서 해제해야 할 일반 목적의 조정은 전혀 가지고 있지 않습니다. 반면에 완성된 체크포인트들은 제가 시도한 모든 학습률에서 거의 움직이지 않거나 심지어 퇴보했습니다.

따라서 여러분이 좋아하는 모델 계열이 사전 지도 학습(pre-supervised) 체크포인트를 게시한다면, 거기서 시작하세요. 그렇지 않다면, 강력한 검색 사전 학습 백본(backbone)에 대한 새로운 투영(projection)이 근접한 차선책입니다. 가장 자연스럽게 느껴지기는 하지만, 도메인 적응을 위해서는 완전히 완성된 체크포인트에서 계속하는 것이 가장 약한 옵션입니다.

MultiVectorEncoderTrainer는 학습 및 평가를 위해 datasets.Dataset 또는 datasets.DatasetDict 인스턴스를 사용합니다. Hugging Face Datasets Hub에서 데이터를 로드하거나 선호하는 형식(예: CSV, JSON, Parquet, Arrow 또는 SQL)의 로컬 데이터를 사용할 수 있습니다.

참고: Sentence Transformers와 함께 바로 사용할 수 있는 많은 공개 데이터셋이 Hugging Face Hub에 sentence-transformers 태그가 지정되어 있습니다. 따라서 https://huggingface.co/datasets?other=sentence-transformers에서 쉽게 찾을 수 있습니다. 이들을 둘러보면서 여러분의 작업, 도메인 또는 언어에 유용할 만한 즉시 사용 가능한 데이터셋을 찾아보세요.

Hub의 데이터셋에서 데이터를 로드하려면 load_dataset 함수를 사용할 수 있습니다:

from datasets import load_dataset
train_dataset = load_dataset(

만약 손실 함수(loss function)가 '손실 개요(Loss Overview)' 표에 따라 *레이블*(Label)을 요구한다면, 데이터셋에는 **"label" 또는 "score"**라는 이름의 컬럼이 있어야 합니다. 이 컬럼은 자동으로 레이블로 간주됩니다. "label"나 "score"라는 이름이 아닌 모든 컬럼은 손실 개요 표에 따라 *입력*(Inputs)으로 간주됩니다. 남아 있는 컬럼의 수는 선택한 손실에 대한 유효 입력의 수와 일치해야 합니다. 이 컬럼들의 이름은 **중요하지 않으며**, 오직 **순서**만 중요합니다.

여기에 두 가지 다중 벡터(multi-vector) 특화 규칙이 추가됩니다:

*   **위치 기반 쿼리 및 문서 할당:** 첫 번째 컬럼은 *쿼리*(query)로 임베딩되고, 그 뒤의 모든 컬럼은 컬럼 이름에 관계없이 *문서*(documents)로 간주됩니다. 이 기본값은 표준 `router_mapping` 훈련 인수를 통해 컬럼별로 재정의할 수 있습니다. 
*   **지식 증류(Knowledge Distillation) 형식:** 각 후보 문서당 하나의 컬럼을 사용합니다. 즉, `(query, document_1, ..., document_N, scores)` 형태이며, 여기서 `scores`는 행당 N개의 교사 점수(teacher scores) 목록입니다. 쿼리와 문서를 별도의 텍스트 데이터셋과 함께 *ID*로 저장하는 KD 데이터셋(예: lightonai/ms-marco-en-bge)의 경우, `resolve_ids`를 사용하여 ID를 실시간으로 텍스트로 변환할 수 있습니다.