인공지능/자연어 처리

Embedding 모델 학습하기 - Sentence Transformer Trainer

이게될까 2025. 9. 7. 01:05
728x90
728x90

https://sbert.net/docs/package_reference/sentence_transformer/trainer.html

 

Trainer — Sentence Transformers documentation

direction (str or List[str], optional, defaults to “minimize”) – If it’s single objective optimization, direction is str, can be “minimize” or “maximize”, you should pick “minimize” when optimizing the validation loss, “maximize” wh

sbert.net

이 사이트에서 많이 참고하면 됩니다.

 

https://huggingface.co/collections/sentence-transformers/embedding-model-datasets-6644d7a3673a511914aa7552

 

Embedding Model Datasets - a sentence-transformers Collection

A curated subset of the datasets that work out of the box with Sentence Transformers: https://huggingface.co/datasets?other=sentence-transformers

huggingface.co

데이터 셋!

 

천천히 작성해 보겠습니다..

 

https://huggingface.co/datasets/sentence-transformers/stsb

 

sentence-transformers/stsb · Datasets at Hugging Face

Dataset Card for STSB The Semantic Textual Similarity Benchmark (Cer et al., 2017) is a collection of sentence pairs drawn from news headlines, video and image captions, and natural language inference data. Each pair is human-annotated with a similarity sc

huggingface.co

sts data set입니다.

문장 2개와 유사도 점수를 가지고 있습니다. 

CosineSimilarityLoss 사용해서 되는 loss!

train 5.75k validation 1.5k test 1.38k 

sts task

 

https://huggingface.co/datasets/sentence-transformers/natural-questions

 

sentence-transformers/natural-questions · Datasets at Hugging Face

Mardi Gras in the United States The expedition, led by Iberville, entered the mouth of the Mississippi River on the evening of March 2, 1699, Lundi Gras, not yet knowing it was the river explored and claimed for France by René-Robert Cavelier, Sieur de La

huggingface.co

NQ Data set!

100k

CachedMultipleNegativesRankingLoss

retrieval task

 

https://huggingface.co/datasets/sentence-transformers/amazon-qa

 

sentence-transformers/amazon-qa · Datasets at Hugging Face

I think it would depend on how you install it. If you have it so the springs loaded in the top hold it in place, then yes you can still use the door. Since it doesn't have a lock once installed, I generally leave a dowel in the track so the slider can't mo

huggingface.co

query answer 

2.51M 

CachedMultipleNegativesRankingLoss

retrieval task

 

https://huggingface.co/datasets/sentence-transformers/gooaq

 

sentence-transformers/gooaq · Datasets at Hugging Face

['Race, Color, or National Origin.', 'Religion.', 'Sex, Gender Identity, or Sexual Orientation.', 'Pregnancy status.', 'Disability.', 'Age or Genetic Information.', 'Citizenship.', 'Marital Status or Number of Children.']

huggingface.co

question answer

3.01M 

CachedMultipleNegativesRankingLoss

retrueval task

 

https://huggingface.co/datasets/bclavie/msmarco-10m-triplets

 

bclavie/msmarco-10m-triplets · Datasets at Hugging Face

Indiana has a wide range of animals. Southern Indiana has some bears and mountain lions, as well as coyotes, white tail deer, and red tail hawks.There are also the usual skunks, raccoons, possums, squirrels, beavers, rabbits, chipmunks and foxes. Indiana h

huggingface.co

query positive negative 

10M 

CachedMultipleNegativesRankingLoss 

retrieval task

 

https://huggingface.co/datasets/sentence-transformers/quora-duplicates

 

sentence-transformers/quora-duplicates · Datasets at Hugging Face

What parallelism can we draw between gods of Hindu mythology, Greek mythology, Egyptian mythology, etc.?

huggingface.co

triplet 102k 

anchor postive negative 

paraphrase

CachedMultipleNegativesRankingLoss 

 

https://huggingface.co/datasets/sentence-transformers/all-nli

 

sentence-transformers/all-nli · Datasets at Hugging Face

Two adults, one female in white, with shades and one male, gray clothes, walking across a street, away from a eatery with a blurred image of a dark colored red shirted person in the foreground.

huggingface.co

triplet 571k

anchor positive negative

nli (관련있는 문장 혹은 유사문장)

CachedMultipleNegativesRankingLoss 

 

https://huggingface.co/datasets/sentence-transformers/stackexchange-duplicates

 

sentence-transformers/stackexchange-duplicates · Datasets at Hugging Face

Let $Q$ be the following set of $\mathbb{Z}\times \mathbb{Z}$ \begin{align*} Q=\{(a,b)\in \mathbb{Z}\times \mathbb{Z} | b\neq 0\} \end{align*} Define the relation $\sim$ on $Q$ as \begin{align*} (a,b)\sim (c,d)\Leftrightarrow ad=bc \end{align*} Prove that

huggingface.co

body-body pair 250k post-post pair 251k title-title pair 305k

body1 body2

retriever

CachedMultipleNegativesRankingLoss 

 

https://huggingface.co/datasets/sentence-transformers/agnews

 

sentence-transformers/agnews · Datasets at Hugging Face

KIRKUK, Iraq - A suicide attacker detonated a car bomb Saturday outside an Iraqi police academy as hundreds of trainees and civilians were leaving for the day, killing 20 people and wounding 36 others in the latest attack designed to thwart U.S-backed effo

huggingface.co

train 1.16M

title description

Retrieval?

CachedMultipleNegativesRankingLoss 

데이터 셋 열(스키마) 샘플 수 task Loss 대안 
sentence-transformers/stsb sentence1, sentence2, score train 5.75k / val 1.5k / test 1.38k STS CosineSimilarityLoss STS에는 회귀형 유사도 손실이 정석. 점수는 0–1 정규화 버전. (Hugging Face)
sentence-transformers/natural-questions question, answer 100,231 검색/질의응답 리트리벌 CachedMultipleNegativesRankingLoss Q–A 페어 기반 인배치 음성 최적. (Hugging Face)
sentence-transformers/amazon-qa query, answer 2,507,114 검색/리트리벌 CachedMultipleNegativesRankingLoss 대규모 Q–A. 캐시형 MNRL로 효율적 학습. (Hugging Face)
sentence-transformers/gooaq question, answer 3,012,496 검색/리트리벌 CachedMultipleNegativesRankingLoss 대규모 Q–A. 페어 스키마 확인. (Hugging Face)
bclavie/msmarco-10m-triplets query, positive, negative 10,000,000 검색/리트리벌 CachedMultipleNegativesRankingLoss 트립렛 채우기에 최적. 페어/트립렛 혼합 전략 가능. (Hugging Face)
sentence-transformers/quora-duplicates (서브셋 다양) pair/pair-class/triplet pair 149k, pair-class 404k, triplet 102k, triplet-all 2.79M 패러프레이즈/중복 질문 CachedMultipleNegativesRankingLoss 중복질문은 MNRL 표준. 서브셋별 크기 확인. (Hugging Face)
sentence-transformers/all-nli 서브셋별 pair-class / pair-score / pair / triplet 전체 2.86M (서브셋 포함) NLI→문장유사/리트리벌 CachedMultipleNegativesRankingLoss pair-score는 CosineSimilarityLoss도 매우 흔함. triplet이면 MNRL/TripletLoss 모두 가능. (Hugging Face)
sentence-transformers/stackexchange-duplicates (서브셋) body1,body2 등 총 805,504 (body-body 250k / post-post 251k / title-title 305k) 중복 질문/패러프레이즈 CachedMultipleNegativesRankingLoss “retriever”라기보다 중복검출/패러프레이즈 쪽에 가까움. (Hugging Face)
sentence-transformers/agnews title, description 1,157,745 뉴스 토픽 분류(4-클래스) (제안: MNRL → 분류/클러스터링용 손실) AGNews는 분류 데이터. 리트리벌로 쓰려면 의도-문서 구조가 아니어서 부자연. SoftmaxLoss(분류) 또는 (Batch)Triplet/Contrastive(클러스터링) 권장. (Hugging Face, Sbert)

 

from datasets import load_dataset
from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer
from sentence_transformers.losses import (
    CachedMultipleNegativesRankingLoss, CosineSimilarityLoss, BatchHardTripletLoss
)
from sentence_transformers.training_args import SentenceTransformerTrainingArguments as STArgs

# 0) 베이스 모델
model = SentenceTransformer("microsoft/mpnet-base")  # 또는 all-MiniLM-L12-v2 등

# 1) 데이터 로딩 (필요량만 샘플링)
msmarco = load_dataset("bclavie/msmarco-10m-triplets", split="train[:2%]")
nq = load_dataset("sentence-transformers/natural-questions", "pair", split="train[:100k]")
gooaq = load_dataset("sentence-transformers/gooaq", "pair", split="train[:200k]")
amazon_qa = load_dataset("sentence-transformers/amazon-qa", "pair", split="train[:200k]")

quora = load_dataset("sentence-transformers/quora-duplicates", split="train[:200k]")
stackdup = load_dataset("sentence-transformers/stackexchange-duplicates", "title-title-pair", split="train[:200k]")

stsb = load_dataset("sentence-transformers/stsb", "pair-score", split="train")
allnli = load_dataset("sentence-transformers/all-nli", "pair-score", split="train[:500k]")

agnews = load_dataset("sentence-transformers/agnews", split="train")  # label 포함 버전
# ※ 컬럼 순서 정리: label/score는 마지막에 두고, 입력 컬럼 2개(or 3개)만 남기기
def ensure_cols(ds, cols): return ds.select_columns(cols)

# Retrieval: (query, positive[, negative])
msmarco = ensure_cols(msmarco, ["query", "positive"])  # 삼중항을 pair로 단순화해 MNRL 사용
nq = nq.rename_columns({"question":"query", "answer":"document"}) if "question" in nq.column_names else nq
gooaq = gooaq.rename_columns({"question":"query", "answer":"document"}) if "question" in gooaq.column_names else gooaq
amazon_qa = amazon_qa.rename_columns({"query":"query", "answer":"document"})

# Paraphrase: (text1, text2)
quora = ensure_cols(quora, ["question1", "question2"])
stackdup = ensure_cols(stackdup, ["body1", "body2"])  # 서브셋에 따라 title-title/body-body 등 선택

# STS/NLI: (text1, text2, label/score)
stsb = ensure_cols(stsb, ["sentence1", "sentence2", "score"])
stsb = stsb.map(lambda e: {"sentence1": e["sentence1"], "sentence2": e["sentence2"], "score": float(e["score"])/5.0})
allnli = ensure_cols(allnli, ["sentence1", "sentence2", "label"])  # 이미 0~1 구간 라벨 제공 서브셋 사용

# 분류/클러스터링: (text, label) → 배치 내 라벨묶음으로 Triplet 구성
agnews = ensure_cols(agnews, ["text", "label"])

# 2) 멀티 데이터셋 딕셔너리
train_datasets = {
    "msmarco": msmarco,
    "nq": nq,
    "gooaq": gooaq,
    "amazon_qa": amazon_qa,
    "quora": quora,
    "stackdup": stackdup,
    "stsb": stsb,
    "allnli": allnli,
    "agnews": agnews,
}

# 3) 데이터셋별 로스
losses = {
    "msmarco": CachedMultipleNegativesRankingLoss(model),
    "nq": CachedMultipleNegativesRankingLoss(model),
    "gooaq": CachedMultipleNegativesRankingLoss(model),
    "amazon_qa": CachedMultipleNegativesRankingLoss(model),
    "quora": CachedMultipleNegativesRankingLoss(model),
    "stackdup": CachedMultipleNegativesRankingLoss(model),
    "stsb": CosineSimilarityLoss(model),     # score ∈ [0,1]
    "allnli": CosineSimilarityLoss(model),   # pair-score 변형
    "agnews": BatchHardTripletLoss(model),   # label 기반
}

# 4) 프롬프트/라우터 매핑
prompts = {
  # 데이터셋별-컬럼별 프롬프트
  "msmarco": {"query": "query: {}", "positive": "passage: {}"},
  "nq": {"query": "question: {}", "document": "answer: {}"},
  "gooaq": {"query": "question: {}", "document": "answer: {}"},
  "amazon_qa": {"query": "question: {}", "document": "answer: {}"},
  # 대칭 과제는 통일 프롬프트(또는 무프롬프트)
  "quora": {"question1": "text: {}", "question2": "text: {}"},
  "stackdup": {"body1": "text: {}", "body2": "text: {}"},
  "stsb": {"sentence1": "text: {}", "sentence2": "text: {}"},
  "allnli": {"sentence1": "text: {}", "sentence2": "text: {}"},
  "agnews": {"text": "text: {}"},
}

router_mapping = {
  "msmarco": {"query": "query", "positive": "document"},
  "nq": {"query": "query", "document": "document"},
  "gooaq": {"query": "query", "document": "document"},
  "amazon_qa": {"query": "query", "document": "document"},
  # 나머지는 대칭 → 동일 라우트(예: 전부 "query" 경로)
  "quora": {"question1": "query", "question2": "query"},
  "stackdup": {"body1": "query", "body2": "query"},
  "stsb": {"sentence1": "query", "sentence2": "query"},
  "allnli": {"sentence1": "query", "sentence2": "query"},
  "agnews": {"text": "query"},
}

# 5) 배치 샘플러: 데이터셋 비율/라벨 묶음/중복 제거
args = STArgs(
    output_dir="exp-multitask",
    per_device_train_batch_size=128,
    learning_rate=2e-5,
    num_train_epochs=1,                         # 또는 max_steps로 제어
    multi_dataset_batch_sampler="proportional", # 데이터셋 크기 비례 샘플
    batch_sampler="no_duplicates+group_by_label",  # MNRL 안정화 + 클러스터링용 라벨 묶음
    prompts=prompts,
    router_mapping=router_mapping,
    evaluation_strategy="steps",
    eval_steps=1000,
    logging_steps=100,
    bf16=True,
    gradient_accumulation_steps=1,
)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_datasets,  # Dict[str, Dataset]
    loss=losses,                   # Dict[str, Loss]
    # evaluators=[...],            # 아래 5) 참조
)

trainer.train()

 

 

 

 

더보기

다음은 SentenceTransformers의 SentenceTransformerTrainer로 임베딩 모델을 학습할 때 필요한 구성요소, 데이터/하이퍼파라미터 설정, 실수하기 쉬운 포인트를 연구자 관점에서 정리한 실전 가이드입니다.

1) 전체 구성요소 개관

SentenceTransformers의 최신 학습 파이프라인은 HF Trainer를 확장한 SentenceTransformerTrainer를 중심으로 돌아갑니다. 핵심 컴포넌트는 6가지입니다: Model, Dataset, Loss, TrainingArguments, Evaluator, Trainer(Callbacks 포함). 예전의 .fit() 방식은 “Deprecated Training”으로 이동했고, 이제는 Trainer 사용이 권장됩니다. (sbert.net)

  • Trainer 핵심: 단일 또는 다중 데이터셋을 받아 학습/평가하며, W&B·TensorBoard·CodeCarbon 등의 콜백을 통합합니다. loss는 하나를 넘겨도 되고, 데이터셋별로 다른 loss를 매핑할 수도 있습니다. loss를 지정하지 않으면 기본적으로 CoSENTLoss가 사용됩니다. (sbert.net)

2) 최소 예시(엔드투엔드)

from datasets import load_dataset
from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer
from sentence_transformers.losses import MultipleNegativesRankingLoss
from sentence_transformers.training_args import SentenceTransformerTrainingArguments
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# 1) 모델
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# 2) 데이터: (anchor, positive) 페어 -> MNRLoss와 호환
train_ds = load_dataset("sentence-transformers/all-nli", "pair", split="train")
dev_ds   = load_dataset("sentence-transformers/all-nli", "pair", split="dev")

# 3) 손실함수
loss = MultipleNegativesRankingLoss(model)

# 4) 트레이닝 아규먼트
args = SentenceTransformerTrainingArguments(
    output_dir="./ckpt",
    do_train=True, do_eval=True,
    evaluation_strategy="steps", eval_steps=1000,
    per_device_train_batch_size=256,  # in-batch negatives ↑
    learning_rate=2e-5, num_train_epochs=1,
    bf16=True, logging_steps=50,
    save_strategy="steps", save_steps=1000,
    load_best_model_at_end=True
)

# 5) 평가자(예: STS 유사도 스피어만)
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=["A man is eating food."],
    sentences2=["A person is eating."],
    scores=[1.0],
    name="dev-sts"
)

# 6) 트레이너
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_ds,
    eval_dataset=None,        # 또는 dev_ds
    loss=loss,
    evaluator=evaluator
)

trainer.train()

위 예시는 (anchor, positive) 쌍을 MultipleNegativesRankingLoss로 학습하는 가장 전형적 형태입니다. 다만 실제 평가는 전용 개발셋/평가자로 교체하세요. 데이터·로스 호환성 원칙은 아래 3) 참고. (sbert.net)

3) 데이터 포맷 규칙(중요)

손실함수와 데이터 포맷이 맞아야 합니다. 규칙은 다음과 같습니다.

  • 라벨이 필요한 loss(예: CosineSimilarityLoss)는 컬럼 이름이 반드시 "label" 또는 "score"여야 합니다.
  • 그 외 컬럼은 입력 텍스트 열로 간주되며, 남은 열의 개수가 해당 loss가 요구하는 입력 개수와 일치해야 합니다(열 이름은 중요하지 않고 열 순서가 중요).
  • 불필요한 메타컬럼은 remove_columns()로 제거하고, 필요시 select_columns()로 열 순서를 재정렬하세요. (sbert.net)

예시 매핑(대표적):

  • MultipleNegativesRankingLoss: ["anchor", "positive"] (라벨 불필요)
  • CosineSimilarityLoss / CoSENTLoss: ["text1", "text2", "label"] (라벨 필요; 유사도 스코어)
  • TripletLoss: ["anchor", "positive", "negative"] (라벨 불필요)
    각 loss의 세부 조건·변형(캐시드 MNRL, 대칭 MNRL, Matryoshka/AdaptiveLayer 등)은 Losses 레퍼런스를 참고하세요. (sbert.net)

4) TrainingArguments 핵심 파라미터(실전 설정 가이드)

SentenceTransformerTrainingArguments는 HF의 TrainingArguments를 확장하여 임베딩 학습에 특화된 인자를 추가합니다. 특히 다음 5가지를 알아두면 좋습니다. (sbert.net)

  1. prompts: 컬럼별/데이터셋별 프롬프트 지정. 예를 들어 비대칭 검색에서
    {"anchor": "query: {text}", "positive": "passage: {text}"} 형태로 문맥 접두사를 일관되게 붙일 수 있습니다(다중 데이터셋이면 데이터셋 이름→컬럼 이름→프롬프트 중첩 딕셔너리). (sbert.net)
  2. router_mapping: 컬럼을 Router 경로(예: "query", "document")에 매핑하여 질의/문서 전용 하위 모듈을 다르게 태운 훈련이 가능합니다(비대칭 검색에 유용). (sbert.net)
  3. learning_rate_mapping: 모듈명 정규식 → LR를 달리 주는 파라미터별 학습률 맵핑(예: 특정 모듈만 더 크게 미세조정). (sbert.net)
  4. 배치/정밀도/스케줄: per_device_train_batch_size, gradient_accumulation_steps, bf16/fp16, warmup_steps/ratio, lr_scheduler_type, save_*, evaluation_strategy, load_best_model_at_end, deepspeed/fsdp/gradient_checkpointing 등은 HF와 동일하게 동작합니다. set_training()/set_dataloader() 헬퍼로 한 번에 세팅 가능. (sbert.net)
  5. 샘플러: batch_sampler, multi_dataset_batch_sampler로 라벨 그룹핑(예: Triplet류)이나 다중 데이터셋 비율 샘플링을 제어합니다. (sbert.net)

5) Loss 선택과 고급 옵션

  • IR/대규모 학습 표준: MultipleNegativesRankingLoss(+ 큰 배치) 추천. 메모리가 빡빡하면 CachedMultipleNegativesRankingLoss 고려(캐시 기반 in-batch negatives 확장). 대칭 태스크면 symmetric 변형 사용 가능. (sbert.net)
  • STS/유사도 회귀: CosineSimilarityLoss 또는 CoSENTLoss. 라벨 포맷 필수. (sbert.net)
  • Triplet 계열: 클래스 라벨 기반일 땐 Triplet류. 배치에 동일 클래스 2개 이상이 들어오도록 샘플러를 맞추는 것이 관건. (sbert.net)
  • 학습 효율/추론 효율 동시 추구:
    • MatryoshkaLoss: 멀티 임베딩 차원에서 동시 최적화(낮은 차원으로도 성능 유지).
    • AdaptiveLayerLoss / Matryoshka2dLoss: 중간 레이어까지 함께 최적화해 “레이어 절단” 시 성능 유지(연구 주제와 잘 맞음). (sbert.net)

6) Evaluator 설정

훈련 중/후 평가는 전용 Evaluator를 사용합니다. 대표적으로:

  • EmbeddingSimilarityEvaluator: STS(스피어만/피어슨).
  • TripletEvaluator / BinaryClassificationEvaluator / InformationRetrievalEvaluator 등 작업별 평가자 제공.
  • 빠른 IR 점검은 NanoBEIR로(정식 BEIR 전 경량 검증).
  • MTEB는 훈련 중 사용을 권장하지 않음(오버피팅 방지 목적) — 훈련 후 외부 벤치 결과 확인 용으로 분리. (sbert.net, sbert.net)

7) 멀티 데이터셋 학습

여러 데이터셋을 서로 다른 포맷/손실로 한 번에 학습할 수 있습니다. train_dataset에 DatasetDict 또는 {name: Dataset} 딕셔너리를 넘기고, loss를 딕셔너리(데이터셋 이름→loss)로 전달하세요. prompts/router_mapping도 데이터셋별로 매핑 가능합니다. (sbert.net)

8) 실무 체크리스트(함정/주의점)

  1. 컬럼 이름·순서 불일치: 라벨 컬럼은 반드시 "label"/"score". 입력 컬럼 개수/순서가 loss 요구사항과 맞는지 확인. 불필요 컬럼은 제거. (sbert.net)
  2. in-batch negatives 오탐: MNRL은 배치 내 다른 샘플을 전부 네거티브로 간주합니다. 중복/유사 문장이 많은 코퍼스는 false negative가 늘 수 있으니, 데이터 클리닝·하드 네거티브 마이닝·캐시드 변형을 고려. (sbert.net)
  3. 비대칭 검색 세팅 누락: 쿼리/문서에 일관된 프롬프트와 router_mapping을 적용하고, 추론 시에도 동일 적용. 학습·추론 프롬프트가 다르면 성능 급락. (sbert.net)
  4. 배치 크기/정밀도: 성능은 종종 “큰 batch(=더 많은 in-batch negatives)”에 비례. GPU 제약 시 bf16/fp16, gradient accumulation, Cached MNRL 사용. (sbert.net)
  5. 평가 빈도/코스트: evaluation_strategy='steps'는 비용이 큽니다. 작은 dev 샘플·느슨한 주기로 시작하고, load_best_model_at_end=True로 early stopping 유사 효과 확보. (sbert.net)
  6. 분산/재현성: deepspeed/fsdp, gradient_checkpointing 사용 시 속도/메모리 이점이 있지만, 완전한 결정론은 보장되지 않습니다. 주요 결과는 여러 시드로 검증. (sbert.net)

9) 자주 쓰는 레시피(권장 기본값)

  • IR(MNRL): lr=2e-5 ~ 5e-5, batch=128~1024(누적 포함), warmup_ratio=0.05, epochs=1~2, bf16=True, evaluation_strategy="steps", eval_steps는 500~2k 정도, save_total_limit=2~3.
  • STS(CoSENT/CosSim): 입력 두 컬럼 + label(0~1 유사도)로 구성, CosineSimilarityEvaluator 병행.
  • 레이어/차원 효율화: AdaptiveLayerLoss(레이어 절단 대비), MatryoshkaLoss(다중 차원)로 효율-성능 균형 설계. (sbert.net)

필요하시면 비대칭 검색(예: query, passage)용 프롬프트/라우터 매핑을 포함한 템플릿이나, 다중 데이터셋(+데이터셋별 loss) 스크립트를 바로 만들어 드리겠습니다. 

 

좋아요. 하나의 임베딩 모델(Shared encoder)로 clustering · STS · retrieval · classification을 동시에 학습하려면, Sentence Transformers의 Multi-Dataset + Multi-Loss 훈련 경로를 쓰는 게 가장 깔끔합니다. 핵심은 아래 5가지입니다:

  1. 데이터를 작업별로 datasets.DatasetDict로 묶고,
  2. 손실함수를 작업별로 다르게 매핑하고,
  3. (비대칭 검색) router_mapping + prompts로 쿼리/문서 라우팅을 분리하고,
  4. multi_dataset_batch_sampler로 작업 믹싱 비율을 통제하고,
  5. 작업별 evaluator로 동시 검증을 돕는 것입니다. (Sbert)

설계 개요

  • 공유 인코더(한 모델): 한 모델 파라미터를 모든 작업이 공유합니다. 필요 시 분기(head)는 훈련 중에만 쓰고(예: SoftmaxLoss), 추론에는 임베딩만 사용합니다. SBERT 자체가 쌍/트리플렛 구조로 임베딩 공간을 직접 최적화하는 철학이므로, 분류 head에 과도하게 의존하지 않게 설계합니다. (arXiv)
  • Multi-dataset / Multi-loss: SentenceTransformerTrainer는 train_dataset=DatasetDict, loss=Dict[str, Loss]를 함께 받으면 데이터셋 이름별로 서로 다른 손실을 적용합니다. (문서에 “loss는 단일 인스턴스 또는 데이터셋 이름→손실 인스턴스 딕셔너리”로 명시) (AiDocZh)
  • Prompt & Router(비대칭 검색): retrieval에서 쿼리/문서 입력을 서로 다른 라우트(예: "query", "document")로 보내도록 router_mapping과, 입력 전처리 프롬프트를 prompts로 작업/컬럼별 지정합니다. 프롬프트·라우터 모두 SentenceTransformerTrainingArguments에서 데이터셋 이름 단위로 매핑 가능합니다. (Sbert)
  • 샘플링 전략: 여러 데이터셋을 섞을 때 multi_dataset_batch_sampler로 "proportional"(기본) 또는 **"round_robin"**을 선택합니다. 전자는 전체 샘플을 다 쓰고 데이터 크기 비례로, 후자는 각 작업을 동일 빈도로 섞되 일부 샘플이 남을 수 있습니다. (Sbert, AiDocZh)
  • 평가: STS에는 EmbeddingSimilarityEvaluator, IR에는 InformationRetrievalEvaluator(또는 NanoBEIR) 등 작업별 evaluator를 함께 걸어 멀티태스킹 중간 성능을 추적합니다. 훈련 후에는 MTEB로 광범위 검증을 권장합니다. (Sbert)

작업별 데이터 형식 & 권장 손실

아래는 Sentence Transformers에서 일반적으로 쓰는 조합입니다. 반드시 데이터 컬럼 형식이 손실이 기대하는 형식과 맞아야 합니다. (문서에서도 “데이터 형식 ↔ 손실함수 호환성”을 강조) (Sbert)

  • STS (문장 유사도 회귀)
    • 데이터: (sentence1, sentence2, score∈[0,1])
    • 손실: CosineSimilarityLoss 혹은 CoSENTLoss (스코어 회귀/순위화) (Sbert)
  • Retrieval (쿼리–문서)
    • 데이터: (query, positive) (+ 가능하면 hard negatives)
    • 손실: MultipleNegativesRankingLoss(MNRL) 또는 대칭 변형(Symmetric) / 대규모면 CachedMultipleNegativesRankingLoss
    • 팁: 라우터(router_mapping)로 query→"query", passage→"document"를 지정, prompts로 “query: … / passage: …” 접두 프롬프트를 부여(비대칭 의미 학습). (AiDocZh, Sbert)
  • Classification (레이블 분류)
    • 옵션 A(임베딩 친화): BatchHardTripletLoss/BatchAllTripletLoss/ContrastiveLoss 등 레이블 기반 metric learning으로 직접 임베딩 공간을 정렬 → 추론에서 임베딩만 써도 성능이 잘 나옵니다. (Sbert)
    • 옵션 B(헤드 부착): SoftmaxLoss(클래스 수 필요). 학습 시 얇은 분류 레이어를 얹습니다. 멀티태스크에서 retrieval/STS와의 간섭을 줄이려면 loss weight나 배치 비율로 영향도를 관리하세요. (Sbert)
  • Clustering (군집)
    • “학습 작업”이라기보다 임베딩 품질 문제입니다. 보통 (레이블이 있으면) Supervised contrastive/Triplet으로 intra-class compactness & inter-class separability를 강화하거나, 레이블이 없으면 SimCSE/CT/TSDAE 같은 자기지도 objective를 추가합니다. (문서의 비지도/증류 레퍼런스 참고) (Sbert)

추가적으로 MatryoshkaLoss / AdaptiveLayerLoss를 “loss modifier”로 감싸 주면, 짧은 벡터나 얕은 레이어로 추론 시에도 성능을 유지하게 만들 수 있습니다(실서비스 속도 최적화에 유용). (Sbert)

최소 동작 예시(청사진)

아래는 “작업별 데이터셋 + 작업별 손실 + 프롬프트/라우터 + 멀티데이터셋 샘플러”의 뼈대입니다. (컬럼명은 여러분 데이터에 맞게 rename_column 하세요.) 샘플링 전략을 바꾸면 멀티태스크 비율을 바로 조정할 수 있습니다. (Hugging Face)

from sentence_transformers import (
    SentenceTransformer, SentenceTransformerTrainer, SentenceTransformerTrainingArguments, losses
)
from datasets import Dataset, DatasetDict

model = SentenceTransformer("microsoft/mpnet-base")  # 베이스는 도메인/언어에 맞게

# 1) 작업별 데이터셋 준비(필요시 rename_column)
train = DatasetDict({
    "sts": Dataset.from_dict({
        "sentence1": s1_list, "sentence2": s2_list, "score": y_float  # [0,1]
    }),
    "retrieval": Dataset.from_dict({
        "query": q_list, "passage": pos_list  # hard negatives는 별도 sampler로
    }),
    "classif": Dataset.from_dict({
        "text": x_list, "label": y_int          # 옵션 A: metric learning 쓰려면 label 유지
    }),
    "cluster": Dataset.from_dict({
        "text": x2_list, "label": y2_int        # triplet/supCon용
    }),
})

# 2) 작업별 손실함수 매핑
loss_map = {
    "sts": losses.CosineSimilarityLoss(model),                 # (sent1, sent2, score)
    "retrieval": losses.MultipleNegativesRankingLoss(model),   # (query, passage)
    # 분류 옵션 A(임베딩 친화)
    "classif": losses.BatchHardTripletLoss(model),
    # 군집도 metric learning로 함께 맞춤
    "cluster": losses.BatchAllTripletLoss(model),
    # ※ 분류 옵션 B를 쓰려면 아래로 교체:
    # "classif": losses.SoftmaxLoss(
    #     model, sentence_embedding_dimension=model.get_sentence_embedding_dimension(),
    #     num_labels=num_classes
    # ),
}

# 3) 훈련 인자: 프롬프트/라우터/샘플러 등
args = SentenceTransformerTrainingArguments(
    output_dir="runs/multitask",
    num_train_epochs=1,
    per_device_train_batch_size=64,
    learning_rate=2e-5,
    evaluation_strategy="steps",
    eval_steps=1000, save_strategy="steps", save_steps=1000, save_total_limit=2,
    multi_dataset_batch_sampler="proportional",  # 또는 "round_robin"
    # 비대칭 검색을 위한 프롬프트 & 라우터
    prompts={
        "retrieval": {"query": "query: ", "passage": "passage: "},
    },
    router_mapping={
        "retrieval": {"query": "query", "passage": "document"},
    },
)

# 4) 트레이너: DatasetDict + dict-of-losses
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train,
    loss=loss_map,          # ← 데이터셋 이름별 손실 적용
    # eval_dataset=eval_dict,  # 선택: 작업별 validation도 DatasetDict로
)

trainer.train()
  • 위 API 패턴은 공식 문서/모듈에 DatasetDict + dict-of-losses 지원이 명시되어 있습니다. (AiDocZh)
  • prompts/router_mapping은 데이터셋 이름→컬럼→설정의 중첩 딕셔너리까지 허용합니다. (Sbert)
  • 멀티데이터셋 샘플러의 정책과 의미는 문서에 정리되어 있습니다. (Sbert)

주의: 손실마다 기대하는 컬럼/레이블 형식이 다릅니다(예: STS는 float score, SoftmaxLoss는 int label, MNRL은 anchor-positive 쌍 등). 손실–데이터 형식을 먼저 맞추세요. (Sbert)

작업 믹싱 & 하이퍼파라미터 팁

  • 비율
    • 검색 성능이 최우선이면 retrieval 배치 비율을 50–70%, 나머지를 STS/분류/클러스터링으로 분할.
    • multi_dataset_batch_sampler="round_robin"은 작업 간 균형을, "proportional"은 전체 데이터 소진을 보장합니다. (Sbert)
  • Similarity 함수 일관성: 모델/손실에 맞춰 cosine vs dot를 통일하세요. 잘못 쓰면 성능이 크게 떨어질 수 있습니다. (MTEB 커뮤니티에서도 주의점으로 언급) (Hugging Face)
  • 프롬프트/라우터: 비대칭 검색에서 “query: … / passage: …” 접두 및 라우팅은 저비용으로 강한 이득을 주는 경우가 많습니다. (공식 예제/인자 지원) (Sbert)
  • Loss modifier: 실서비스 속도·메모리 제약이 있다면 MatryoshkaLoss(차원 절단 내성), AdaptiveLayerLoss(얕은 레이어 추론 내성)를 outer loss로 감싸는 것을 권장. (Sbert)
  • Hard negative 관리: MNRL은 배치 내 **거짓 음성(false negative)**에 민감합니다. 도메인 전처리/약한 필터, 오프라인 hard-negative 채굴, 혹은 util의 hard negative 유틸을 활용하세요. (Sbert)
  • 평가 루프: 훈련 중에는 STS(EmbeddingSimilarityEvaluator), IR(InformationRetrievalEvaluator/NanoBEIR)을 주기적으로 돌리고, 최종은 MTEB로 종합평가하세요. (Sbert)
  • PEFT/Adapter: 간섭을 줄이거나 메모리를 아끼려면 작업별 어댑터(PEFT)로 파라미터를 분리해보는 것도 실용적입니다. (공식 예제 제공) (Sbert)

실무 체크리스트 (함정/주의)

  • (가장 중요) 손실–데이터 호환성 확인(컬럼명/레이블 타입/스케일). (Sbert)
  • 라벨 스케일: STS score는 [0,1]이 안전합니다(데이터셋마다 0–5인 경우 normalization 필요).
  • 유사도 함수: 훈련/평가/서빙에서 동일한 유사도(cosine or dot) 사용. (Hugging Face)
  • 작업 간 간섭: SoftmaxLoss를 섞을 때 retrieval가 흔들릴 수 있습니다. → 분류는 가급적 metric-learning 계열로 시작하거나, round_robin + 낮은 loss weight로 완화. (Sbert)
  • 샘플 불균형: 데이터셋 크기가 크게 다르면 "proportional"에서 큰 데이터가 학습을 지배합니다. 필요한 경우 "round_robin"으로 공정 혼합. (Sbert)
  • 로깅/평가지표: 작업별 evaluator 지표가 다릅니다(STS=Spearman, IR=MRR/nDCG, 분류=ACC/F1). 문서의 evaluator 표를 기준으로 설정하세요. (Sbert)

한 줄 권장 레시피

Retrieval 중심 + 범용 임베딩이 목표라면
MNRL(주), STS 보조(CoSENT/Cosine), 분류·클러스터링은 Triplet/Contrastive로 묶고, 프롬프트/라우터를 켠 뒤 **round_robin**으로 작업 균형을 잡는 구성을 1차 권장합니다. 이후 Matryoshka/AdaptiveLayer로 “짧은 벡터/얕은 레이어 추론” 최적화를 추가하세요. (AiDocZh, Sbert)

필요하시면, 실제로 쓰실 데이터셋(컬럼명/크기/언어/우선순위)에 맞춰 손실 조합·배치 비율·프롬프트/라우팅·평가 설정을 바로 짜 드릴게요.

 

아래는 Hugging Face의 Embedding Model Datasets 컬렉션에서 과제를 고루 커버하는 데이터를 취합해, Sentence Transformers(v3)로 retrieval / STS / paraphrase(dup) / classification / clustering를 한 모델로 멀티태스킹 학습하는 실전 레시피입니다. 각 단계는 SBERT v3의 SentenceTransformerTrainer/SentenceTransformerTrainingArguments가 제공하는 다중 데이터셋·다중 로스 파이프라인을 전제로 합니다. (Hugging Face, Sbert)

1) 데이터 구성(추천 조합)

컬렉션에서 바로 쓸 수 있고 스키마가 단순한 것들 위주로 추렸습니다(괄호는 추천 용도).

Retrieval(비대칭, 쿼리→문서)

  • MS MARCO: bclavie/msmarco-10m-triplets (query, positive, negative) — 대규모 삼중항. 삼중항을 pair로도 사용 가능(MNRL의 in-batch negative). (Hugging Face)
  • Natural Questions (pair): (query, answer). (Hugging Face)
  • GooAQ (pair): (question, answer). (Hugging Face)
  • Amazon Q&A (pair): (query, answer). (Hugging Face)

Paraphrase / Duplicate(대칭)

  • Quora duplicates (pair), StackExchange duplicates (pair: title-title / body-body / post-post). (Hugging Face)

STS / NLI(연속 레이블)

  • STS-B (pair-score) — label(유사도) ∈ [0, 5] → [0,1] 정규화 필요.
  • AllNLI (pair-score / 또는 pair + 3-class) — NLI를 유사도 점수(contradiction=0, neutral≈0.5, entailment=1)로 매핑해 회귀 계열 로스 사용. (Hugging Face)

분류/클러스터링(레이블 기반)

  • AG News(label, title/description) — 동일 라벨을 클러스터로 간주해 Triplet류 사용.
  • Amazon Reviews(title, review) — 별점 등 라벨이 있는 버전을 활용하면 분류/클러스터링에 유용. (Hugging Face)

이 컬렉션들은 SBERT v3에서 “바로 학습 가능한” 형식으로 큐레이션되어 있어, 최소한의 컬럼 정리만으로 곧장 쓸 수 있습니다. (Hugging Face)

2) 손실 함수(로스) 매핑

  • Retrieval (pair: query, doc) → CachedMultipleNegativesRankingLoss (대규모에서 안정적), 또는 GISTEmbedLoss(가이드 모델을 이용해 더 강한 신호, 메모리-시간 트레이드오프). (Sbert)
  • Paraphrase / Duplicate (pair) → 위와 동일하게 CachedMultipleNegativesRankingLoss(대칭).
  • STS / NLI (pair + score/label) → CosineSimilarityLoss 또는 CoSENTLoss 같은 점수 회귀형 로스. 데이터 포맷 규칙: 레이블 컬럼은 반드시 "label" 또는 "score"이고, 나머지 입력 컬럼 수가 해당 로스의 입력 수와 일치해야 합니다(컬럼 이름은 무관, 순서가 중요). (Sbert)
  • Classification / Clustering (label 기반) → BatchHardTripletLoss 등 Triplet류(동일 라벨을 양성으로), SBERT v3의 batch sampler에서 group_by_label을 쓰면 배치 내에 동일 라벨 쌍을 보장할 수 있습니다.

3) 비대칭 검색을 위한 Router & Prompt

SBERT v3는 Router + prompts로 쿼리/문서 경로를 분리합니다. TrainingArguments의

  • prompts: 컬럼 또는 데이터셋별 프롬프트 문자열 매핑
  • router_mapping: (dataset→){column→"query"/"document"} 경로 매핑
  • encode_query() / encode_document() 사용 권장(추론 시 자동으로 적절한 프롬프트/라우터를 사용)
    을 지원합니다. Retrieval 성능은 간단한 프롬프트만으로도 유의미하게 상승하는 경우가 보고되어 있습니다. (Sbert)

4) 멀티데이터셋·멀티로스 트레이너 스켈레톤

from datasets import load_dataset
from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer
from sentence_transformers.losses import (
    CachedMultipleNegativesRankingLoss, CosineSimilarityLoss, BatchHardTripletLoss
)
from sentence_transformers.training_args import SentenceTransformerTrainingArguments as STArgs

# 0) 베이스 모델
model = SentenceTransformer("microsoft/mpnet-base")  # 또는 all-MiniLM-L12-v2 등

# 1) 데이터 로딩 (필요량만 샘플링)
msmarco = load_dataset("bclavie/msmarco-10m-triplets", split="train[:2%]")
nq = load_dataset("sentence-transformers/natural-questions", "pair", split="train[:100k]")
gooaq = load_dataset("sentence-transformers/gooaq", "pair", split="train[:200k]")
amazon_qa = load_dataset("sentence-transformers/amazon-qa", "pair", split="train[:200k]")

quora = load_dataset("sentence-transformers/quora-duplicates", split="train[:200k]")
stackdup = load_dataset("sentence-transformers/stackexchange-duplicates", "title-title-pair", split="train[:200k]")

stsb = load_dataset("sentence-transformers/stsb", "pair-score", split="train")
allnli = load_dataset("sentence-transformers/all-nli", "pair-score", split="train[:500k]")

agnews = load_dataset("sentence-transformers/agnews", split="train")  # label 포함 버전
# ※ 컬럼 순서 정리: label/score는 마지막에 두고, 입력 컬럼 2개(or 3개)만 남기기
def ensure_cols(ds, cols): return ds.select_columns(cols)

# Retrieval: (query, positive[, negative])
msmarco = ensure_cols(msmarco, ["query", "positive"])  # 삼중항을 pair로 단순화해 MNRL 사용
nq = nq.rename_columns({"question":"query", "answer":"document"}) if "question" in nq.column_names else nq
gooaq = gooaq.rename_columns({"question":"query", "answer":"document"}) if "question" in gooaq.column_names else gooaq
amazon_qa = amazon_qa.rename_columns({"query":"query", "answer":"document"})

# Paraphrase: (text1, text2)
quora = ensure_cols(quora, ["question1", "question2"])
stackdup = ensure_cols(stackdup, ["body1", "body2"])  # 서브셋에 따라 title-title/body-body 등 선택

# STS/NLI: (text1, text2, label/score)
stsb = ensure_cols(stsb, ["sentence1", "sentence2", "score"])
stsb = stsb.map(lambda e: {"sentence1": e["sentence1"], "sentence2": e["sentence2"], "score": float(e["score"])/5.0})
allnli = ensure_cols(allnli, ["sentence1", "sentence2", "label"])  # 이미 0~1 구간 라벨 제공 서브셋 사용

# 분류/클러스터링: (text, label) → 배치 내 라벨묶음으로 Triplet 구성
agnews = ensure_cols(agnews, ["text", "label"])

# 2) 멀티 데이터셋 딕셔너리
train_datasets = {
    "msmarco": msmarco,
    "nq": nq,
    "gooaq": gooaq,
    "amazon_qa": amazon_qa,
    "quora": quora,
    "stackdup": stackdup,
    "stsb": stsb,
    "allnli": allnli,
    "agnews": agnews,
}

# 3) 데이터셋별 로스
losses = {
    "msmarco": CachedMultipleNegativesRankingLoss(model),
    "nq": CachedMultipleNegativesRankingLoss(model),
    "gooaq": CachedMultipleNegativesRankingLoss(model),
    "amazon_qa": CachedMultipleNegativesRankingLoss(model),
    "quora": CachedMultipleNegativesRankingLoss(model),
    "stackdup": CachedMultipleNegativesRankingLoss(model),
    "stsb": CosineSimilarityLoss(model),     # score ∈ [0,1]
    "allnli": CosineSimilarityLoss(model),   # pair-score 변형
    "agnews": BatchHardTripletLoss(model),   # label 기반
}

# 4) 프롬프트/라우터 매핑
prompts = {
  # 데이터셋별-컬럼별 프롬프트
  "msmarco": {"query": "query: {}", "positive": "passage: {}"},
  "nq": {"query": "question: {}", "document": "answer: {}"},
  "gooaq": {"query": "question: {}", "document": "answer: {}"},
  "amazon_qa": {"query": "question: {}", "document": "answer: {}"},
  # 대칭 과제는 통일 프롬프트(또는 무프롬프트)
  "quora": {"question1": "text: {}", "question2": "text: {}"},
  "stackdup": {"body1": "text: {}", "body2": "text: {}"},
  "stsb": {"sentence1": "text: {}", "sentence2": "text: {}"},
  "allnli": {"sentence1": "text: {}", "sentence2": "text: {}"},
  "agnews": {"text": "text: {}"},
}

router_mapping = {
  "msmarco": {"query": "query", "positive": "document"},
  "nq": {"query": "query", "document": "document"},
  "gooaq": {"query": "query", "document": "document"},
  "amazon_qa": {"query": "query", "document": "document"},
  # 나머지는 대칭 → 동일 라우트(예: 전부 "query" 경로)
  "quora": {"question1": "query", "question2": "query"},
  "stackdup": {"body1": "query", "body2": "query"},
  "stsb": {"sentence1": "query", "sentence2": "query"},
  "allnli": {"sentence1": "query", "sentence2": "query"},
  "agnews": {"text": "query"},
}

# 5) 배치 샘플러: 데이터셋 비율/라벨 묶음/중복 제거
args = STArgs(
    output_dir="exp-multitask",
    per_device_train_batch_size=128,
    learning_rate=2e-5,
    num_train_epochs=1,                         # 또는 max_steps로 제어
    multi_dataset_batch_sampler="proportional", # 데이터셋 크기 비례 샘플
    batch_sampler="no_duplicates+group_by_label",  # MNRL 안정화 + 클러스터링용 라벨 묶음
    prompts=prompts,
    router_mapping=router_mapping,
    evaluation_strategy="steps",
    eval_steps=1000,
    logging_steps=100,
    bf16=True,
    gradient_accumulation_steps=1,
)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_datasets,  # Dict[str, Dataset]
    loss=losses,                   # Dict[str, Loss]
    # evaluators=[...],            # 아래 5) 참조
)

trainer.train()
  • 위처럼 데이터셋 딕셔너리와 로스 딕셔너리를 주면, 트레이너가 데이터셋별 맞는 로스를 적용합니다.
  • prompts/router_mapping은 데이터셋·컬럼별로 지정 가능합니다.
  • multi_dataset_batch_sampler="proportional"(기본)은 데이터셋 크기에 비례해 배치를 샘플링, round_robin은 균등 순환을 만듭니다.
  • batch_sampler에 no_duplicates(MNRL 학습 안정화)나 group_by_label(TripletLoss용)을 합성해 사용할 수 있습니다. (Sbert)

데이터셋-로스 호환성 체크: SBERT v3는 “레이블 컬럼은 label/score여야 하고, 나머지 입력 컬럼 수가 로스의 입력 수와 일치”해야 함을 명시합니다. 필요시 select_columns/remove_columns/rename_columns로 정리하세요(컬럼명 자체는 중요치 않지만 순서는 중요). (Sbert)

5) 중간/최종 평가(Eval)

  • 학습 중 빠른 IR 점검: NanoBEIR(Evaluator) 또는 InformationRetrievalEvaluator를 소규모 코퍼스로 구성해 evaluation_strategy="steps"로 주기 평가. 프롬프트 사용 시 IR이 소폭 개선된 예제가 공식 스크립트로 제공됩니다. (Sbert)
  • STS: EmbeddingSimilarityEvaluator(Spearman/Pearson)로 STS-dev를 정기 측정. (Sbert)
  • 파라프레이즈: BinaryClassificationEvaluator로 dup/non-dup AUC/F1. (Sbert)
  • 최종 총괄: 학습 완료 후 MTEB로 전범위(8개 카테고리) 평가—리더보드 제출도 가능. (훈련 중 공개 벤치에 과적합하지 않도록 최종에 사용 권장) (Sbert, GitHub)

6) 하이퍼파라미터·샘플링 가이드(실전값)

  • 배치: GPU 메모리에 맞추어 64–256. MNRL은 큰 배치가 이득.
  • 학습률: 1e-5 ~ 3e-5(대형/사전학습 강한 베이스는 더 작게).
  • 스텝 비율(대략): Retrieval 60–70%, Paraphrase 10–15%, STS/NLI 10–15%, 분류/클러스터링 10–15%. multi_dataset_batch_sampler로 비율 조정. (Sbert)
  • 프롬프트: "query: {} vs "passage: {}"처럼 단순 프롬프트만으로도 NDCG@10이 개선된 사례가 공식 예제에 포함. (Sbert)
  • 추론: IR에는 encode_query() / encode_document()를 사용해 라우터·프롬프트 자동 적용. (Sbert)

7) 특히 주의할 점 (실패 패턴 방지 체크리스트)

  1. 데이터 포맷 불일치
    • 로스가 요구하는 입력 수/순서와 데이터 컬럼 순서를 반드시 맞추세요. label/score 이름은 고정이어야 합니다. (Sbert)
  2. Retrieval 비대칭성 무시
    • 쿼리/문서를 동일 경로로 학습하면 성능 저하. router_mapping과 prompts를 세팅하고, 추론 시 encode_query/document()를 쓰세요. (Sbert)
  3. In-batch negative 오염(중복 텍스트)
    • MNRL에서 같은 문장이 한 배치에 중복되면 collapse 위험. batch_sampler="no_duplicates"를 권장.
  4. 레이블 기반 로스의 배치 구성
    • Triplet류는 배치 내 동일 라벨 샘플이 충분히 있어야 합니다. group_by_label를 켜고, 라벨 불균형이면 클래스별 샘플링 가중치를 고려.
  5. STS 점수 스케일
    • STS-B는 0–5 스케일 → [0,1]로 정규화 후 회귀 로스 사용. (공식 STS 예제·도큐에서 권장) (Sbert)
  6. 데이터 편중
    • MS MARCO 등 초대규모가 전체를 압도하지 않게 샘플 비율을 조정하거나 round_robin 샘플러를 고려. (Sbert)
  7. 평가 데이터 누수/벤치마크 과적합
    • MTEB는 학습 종료 후에만 사용하고, 개발 중에는 NanoBEIR/STS-dev 같은 소규모로 트렌드만 모니터링. (Sbert)
  8. 프롬프트·길이 제한
    • 프롬프트와 최대 시퀀스 길이(예: 512)로 인한 토큰 예산을 점검. IR은 쿼리/문서에 서로 다른 프롬프트를 적용.
  9. 멀티링구얼
    • 다국어가 목표면 parallel-sentences-*, MIRACL/Mr.TyDi 등 멀티링구얼 세트를 추가하고, 라우터·프롬프트를 언어마다 일관되게. (Hugging Face)
  10. Matryoshka(선택)
  • 하나의 모델에서 여러 임베딩 차원을 동시에 최적화해 추론비용을 줄일 수 있음(SBERT 예제/블로그 제공). 대규모 프로덕션에 유리. (Sbert, Hugging Face)

8) 최소 예시: 평가기 추가

from sentence_transformers import evaluation

# STS 유사도 상관
dev_sts = load_dataset("sentence-transformers/stsb", "pair-score", split="validation")
dev_sts = dev_sts.select_columns(["sentence1", "sentence2", "score"])
dev_sts = dev_sts.map(lambda e: {"score": float(e["score"])/5.0})
e_sts = evaluation.EmbeddingSimilarityEvaluator(
    dev_sts["sentence1"], dev_sts["sentence2"], dev_sts["score"], name="sts-dev"
)

# 간단 IR 평가 (NanoBEIR/IR Evaluator를 small set으로 구성)
# ... (질의/코퍼스/레버런스 준비 후) e_ir = evaluation.InformationRetrievalEvaluator(...)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_datasets,
    loss=losses,
    evaluators=[e_sts],  # , e_ir
)
trainer.train()

평가기 종류와 사용법은 evaluation 레퍼런스 및 “Training with prompts” 예제( NanoBEIR 사용 )에 정리되어 있습니다. (Sbert)

마무리: 실행 순서 요약

  1. 베이스 모델 선택(mpnet/miniLM 등)
  2. 데이터셋 로딩·정규화(컬럼 순서/레이블 스케일/프롬프트·라우팅 설정)
  3. 과제별 로스 구성 + 배치 샘플러(no_duplicates, group_by_label, proportional/round_robin)
  4. 중간평가(STS-dev, NanoBEIR)로 드리프트 감시
  5. 최종 MTEB 평가로 범용성 검증 및 리더보드 제출(선택) (Sbert)

필요하시면, 위 스켈레톤을 **당신의 도메인/언어(KO 등)**로 구체화하여 정확한 데이터 비율·하이퍼파라미터·프롬프트 설계까지 포함한 실행 스크립트로 바로 드리겠습니다.

 

 

더보기

Multi-task Embedding Model 훈련 가이드

embedding 모델을 STS, Clustering, Summarization, Retrieval, Classification 등 모든 task에서 잘 작동하도록 학습시키는 완전한 가이드를 제공해드리겠습니다.

🎯 Task별 데이터 및 Loss Function

1. STS (Semantic Textual Similarity)

·         데이터 형식: (sentence1, sentence2, similarity_score)

·         데이터셋: STS-B, KorSTS, SICK-R

·         Loss: CoSENTLoss 또는 MSELoss

·         설명: 두 문장 간 의미적 유사도를 0-1 점수로 학습

2. NLI (Natural Language Inference)

·         데이터 형식: (premise, hypothesis, label)

·         데이터셋: SNLI, MultiNLI, KorNLI

·         Loss: SoftmaxLoss (3-class)

·         설명: 전제와 가설 간 논리적 관계 분류

3. Retrieval

·         데이터 형식: (query, positive_doc, negative_doc)

·         데이터셋: MS-MARCO, Natural Questions, HotpotQA

·         Loss: MultipleNegativesRankingLoss (InfoNCE)

·         설명: 쿼리에 대한 관련 문서 검색

4. Classification

·         데이터 형식: (text, class_label)

·         데이터셋: AG News, Amazon Reviews, 한국어 리뷰 데이터

·         Loss: SoftmaxLoss

·         설명: 텍스트를 사전정의된 카테고리로 분류

5. Clustering

·         데이터 형식: (text, cluster_label)

·         데이터셋: 20 Newsgroups, ArXiv papers

·         Loss: InfoNCE with cluster-based sampling

·         설명: 유사한 텍스트들을 그룹화

🚀 실제 훈련 코드

기본적인 multi-task 훈련 코드는 다음과 같습니다:

import torch
from datasets import load_dataset, Dataset
from sentence_transformers import (
    SentenceTransformer, 
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments
)
from sentence_transformers.losses import (
    MultipleNegativesRankingLoss,
    CoSENTLoss,
    SoftmaxLoss
)
from sentence_transformers.evaluation import (
    EmbeddingSimilarityEvaluator,
    TripletEvaluator,
    InformationRetrievalEvaluator
)
from sentence_transformers.training_args import BatchSamplers

# 1. 모델 초기화
model = SentenceTransformer("klue/bert-base")  # 한국어 베이스 모델 사용
print(f"Model embedding dimension: {model.get_sentence_embedding_dimension()}")

# 2. 다양한 태스크별 데이터셋 로드
def load_multitask_datasets():
    """다양한 태스크의 데이터셋을 로드하고 포맷을 맞춤"""

    datasets = {}

    # STS 데이터 (sentence1, sentence2, score)
    sts_data = load_dataset("sentence-transformers/stsb", split="train[:10000]")
    datasets["sts"] = sts_data

    # NLI 데이터 (premise, hypothesis, label) -> (sentence1, sentence2, class)
    nli_data = load_dataset("sentence-transformers/all-nli", "pair-class", split="train[:10000]")
    datasets["nli"] = nli_data

    # Retrieval 데이터 (anchor, positive, negative)
    retrieval_data = load_dataset("sentence-transformers/all-nli", "triplet", split="train[:10000]")
    datasets["retrieval"] = retrieval_data

    # Paraphrase 데이터 (anchor, positive)
    paraphrase_data = load_dataset("sentence-transformers/quora-duplicates", "pair", split="train[:10000]")
    datasets["paraphrase"] = paraphrase_data

    return datasets

# 3. 각 태스크별 Loss Function 정의
def setup_loss_functions(model):
    """태스크별 Loss Function 초기화"""

    losses = {
        # STS: 연속적인 유사도 점수 학습
        "sts": CoSENTLoss(model),

        # NLI: 3-class 분류 (entailment, neutral, contradiction)
        "nli": SoftmaxLoss(
            model=model,
            sentence_embedding_dimension=model.get_sentence_embedding_dimension(),
            num_labels=3
        ),

        # Retrieval: Contrastive Learning with InfoNCE
        "retrieval": MultipleNegativesRankingLoss(model),

        # Paraphrase: Contrastive Learning
        "paraphrase": MultipleNegativesRankingLoss(model)
    }

    return losses

# 4. 평가자(Evaluator) 설정
def setup_evaluators():
    """각 태스크별 평가자 설정"""

    # STS 평가용 데이터 로드
    sts_eval = load_dataset("sentence-transformers/stsb", split="validation[:1000]")

    evaluators = []

    # STS Evaluator
    sts_evaluator = EmbeddingSimilarityEvaluator(
        sentences1=sts_eval["sentence1"],
        sentences2=sts_eval["sentence2"], 
        scores=sts_eval["score"],
        name="sts-eval"
    )
    evaluators.append(sts_evaluator)

    # Triplet Evaluator for Retrieval
    triplet_eval = load_dataset("sentence-transformers/all-nli", "triplet", split="dev[:1000]")
    triplet_evaluator = TripletEvaluator(
        anchors=triplet_eval["anchor"],
        positives=triplet_eval["positive"],
        negatives=triplet_eval["negative"],
        name="triplet-eval"
    )
    evaluators.append(triplet_evaluator)

    return evaluators

# 5. Training Arguments 설정
def setup_training_args():
    """훈련 파라미터 설정"""

    args = SentenceTransformerTrainingArguments(
        output_dir="./multi-task-embedding-model",

        # 기본 훈련 설정
        num_train_epochs=3,
        per_device_train_batch_size=32,
        per_device_eval_batch_size=32,

        # 학습률 및 스케줄러
        learning_rate=2e-5,
        warmup_ratio=0.1,
        lr_scheduler_type="cosine",

        # Mixed Precision Training
        fp16=True,  # GPU 메모리 절약 및 속도 향상

        # Multi-dataset 배치 샘플링
        batch_sampler=BatchSamplers.NO_DUPLICATES,  # InfoNCE에 유리
        multi_dataset_batch_sampler="proportional",  # 데이터셋 크기 비례 샘플링

        # 평가 설정
        eval_strategy="steps",
        eval_steps=500,
        save_strategy="steps", 
        save_steps=500,
        save_total_limit=2,

        # 로깅
        logging_steps=100,
        report_to="wandb",  # Weights & Biases 로깅
        run_name="multi-task-embedding-klue-bert",

        # 최적화
        load_best_model_at_end=True,
        metric_for_best_model="eval_sts-eval_spearman_cosine"
    )

    return args

# 6. Hard Negative Mining 함수
def mine_hard_negatives(model, queries, passages, top_k=50):
    """하드 네거티브 샘플 마이닝"""

    import faiss
    import numpy as np

    # 문서들을 임베딩
    doc_embeddings = model.encode(passages, convert_to_tensor=False)
    query_embeddings = model.encode(queries, convert_to_tensor=False)

    # FAISS 인덱스 구축
    dimension = doc_embeddings.shape[1]
    index = faiss.IndexFlatIP(dimension)  # Inner Product (cosine similarity)

    # 정규화 후 인덱스에 추가
    doc_embeddings = doc_embeddings / np.linalg.norm(doc_embeddings, axis=1, keepdims=True)
    index.add(doc_embeddings.astype('float32'))

    # 각 쿼리에 대해 유사한 문서 검색
    query_embeddings = query_embeddings / np.linalg.norm(query_embeddings, axis=1, keepdims=True)
    scores, indices = index.search(query_embeddings.astype('float32'), top_k)

    return indices, scores

# 7. 메인 훈련 함수
def train_multitask_embedding_model():
    """Multi-task 임베딩 모델 훈련"""

    # 모델 초기화
    model = SentenceTransformer("klue/bert-base")

    # 데이터셋 로드
    print("Loading datasets...")
    train_datasets = load_multitask_datasets()

    # Loss functions 설정
    print("Setting up loss functions...")
    loss_functions = setup_loss_functions(model)

    # 평가자 설정
    print("Setting up evaluators...")
    evaluators = setup_evaluators()

    # 훈련 파라미터 설정
    print("Setting up training arguments...")
    args = setup_training_args()

    # Trainer 생성
    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_datasets,
        loss=loss_functions,
        evaluator=evaluators
    )

    # 훈련 시작
    print("Starting training...")
    trainer.train()

    # 모델 저장
    model.save_pretrained("./multi-task-embedding-final")

    return model, trainer

# 8. 모델 평가 함수
def evaluate_model(model):
    """훈련된 모델의 성능 평가"""

    # 예시 문장들
    sentences = [
        "안녕하세요, 반갑습니다.",
        "Hello, nice to meet you.",
        "오늘 날씨가 좋네요.",
        "The weather is nice today.",
        "인공지능은 미래 기술입니다.",
        "AI is the technology of the future."
    ]

    # 임베딩 생성
    embeddings = model.encode(sentences)

    # 유사도 계산
    from sentence_transformers.util import cos_sim
    similarities = cos_sim(embeddings, embeddings)

    print("\nSentence Similarities:")
    for i, sent1 in enumerate(sentences):
        for j, sent2 in enumerate(sentences):
            if i < j:
                sim = similarities[i][j].item()
                print(f"'{sent1}' vs '{sent2}': {sim:.4f}")

    return embeddings, similarities

# 사용 예시
if __name__ == "__main__":
    # 훈련 실행
    model, trainer = train_multitask_embedding_model()

    # 모델 평가
    embeddings, similarities = evaluate_model(model)

    print("Multi-task embedding model training completed!")

고급 기능을 포함한 완전한 훈련 파이프라인:

import torch
import numpy as np
import pandas as pd
from typing import List, Dict, Tuple, Optional
import json
import random
from datasets import Dataset, DatasetDict
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer
import faiss

class HardNegativeMiner:
    """하드 네거티브 샘플 마이닝 클래스"""

    def __init__(self, model: SentenceTransformer, top_k: int = 50):
        self.model = model
        self.top_k = top_k
        self.index = None
        self.passages = None

    def build_index(self, passages: List[str]):
        """문서 컬렉션으로부터 FAISS 인덱스 구축"""
        print(f"Building FAISS index for {len(passages)} passages...")

        # 문서 임베딩 생성
        embeddings = self.model.encode(
            passages, 
            batch_size=32, 
            show_progress_bar=True,
            convert_to_tensor=False
        )

        # 정규화
        embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)

        # FAISS 인덱스 생성
        dimension = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dimension)
        self.index.add(embeddings.astype('float32'))

        self.passages = passages
        print(f"Index built successfully with {self.index.ntotal} vectors")

    def mine_negatives(self, queries: List[str], positive_indices: List[int] = None) -> List[List[str]]:
        """쿼리에 대한 하드 네거티브 마이닝"""
        if self.index is None:
            raise ValueError("Index not built. Call build_index() first.")

        print(f"Mining hard negatives for {len(queries)} queries...")

        # 쿼리 임베딩 생성
        query_embeddings = self.model.encode(queries, convert_to_tensor=False)
        query_embeddings = query_embeddings / np.linalg.norm(query_embeddings, axis=1, keepdims=True)

        # 검색 실행
        scores, indices = self.index.search(query_embeddings.astype('float32'), self.top_k)

        hard_negatives = []
        for i, (query_indices, query_scores) in enumerate(zip(indices, scores)):
            negatives = []
            positive_idx = positive_indices[i] if positive_indices else None

            for idx, score in zip(query_indices, query_scores):
                # positive sample 제외
                if positive_idx is not None and idx == positive_idx:
                    continue

                # 적절한 유사도 범위의 하드 네거티브만 선택
                if 0.3 < score < 0.8:  # 너무 쉽거나 너무 어려운 것 제외
                    negatives.append(self.passages[idx])

                if len(negatives) >= 5:  # 네거티브 개수 제한
                    break

            hard_negatives.append(negatives)

        return hard_negatives

class MultiTaskDataProcessor:
    """다중 태스크 데이터 전처리 클래스"""

    def __init__(self, tokenizer_name: str = "klue/bert-base"):
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        self.max_length = 512

    def process_sts_data(self, data: List[Dict]) -> Dataset:
        """STS 데이터 전처리"""
        processed = []

        for item in data:
            processed.append({
                'sentence1': item['sentence1'],
                'sentence2': item['sentence2'],
                'score': float(item['score']) / 5.0  # 0-5 스케일을 0-1로 정규화
            })

        return Dataset.from_list(processed)

    def process_nli_data(self, data: List[Dict]) -> Dataset:
        """NLI 데이터 전처리"""
        processed = []
        label_map = {'entailment': 0, 'neutral': 1, 'contradiction': 2}

        for item in data:
            if item['label'] in label_map:
                processed.append({
                    'premise': item['premise'],
                    'hypothesis': item['hypothesis'],
                    'label': label_map[item['label']]
                })

        return Dataset.from_list(processed)

    def process_retrieval_data(self, queries: List[str], documents: List[str], 
                             relevant_docs: List[List[int]], hard_negatives: List[List[str]] = None) -> Dataset:
        """검색 데이터 전처리 (하드 네거티브 포함)"""
        processed = []

        for i, query in enumerate(queries):
            relevant_indices = relevant_docs[i]
            query_hard_negs = hard_negatives[i] if hard_negatives else []

            for rel_idx in relevant_indices:
                # Positive pair
                processed.append({
                    'anchor': query,
                    'positive': documents[rel_idx]
                })

                # Hard negatives 추가
                for hard_neg in query_hard_negs[:3]:  # 최대 3개
                    processed.append({
                        'anchor': query,
                        'positive': documents[rel_idx],
                        'negative': hard_neg
                    })

        return Dataset.from_list(processed)

    def create_instruction_data(self, text: str, task_type: str) -> str:
        """Instruction-following 형태의 데이터 생성 (E5/BGE 스타일)"""
        instructions = {
            'retrieval': "Represent this sentence for searching relevant passages: ",
            'classification': "Represent this sentence for classification: ",
            'clustering': "Represent this sentence for clustering: ",
            'sts': "Represent this sentence for measuring semantic similarity: "
        }

        instruction = instructions.get(task_type, "")
        return instruction + text

    def balance_dataset(self, datasets: Dict[str, Dataset], max_samples: int = 50000) -> Dict[str, Dataset]:
        """데이터셋 균형 맞추기"""
        balanced = {}

        for name, dataset in datasets.items():
            if len(dataset) > max_samples:
                # 랜덤 샘플링
                indices = random.sample(range(len(dataset)), max_samples)
                balanced[name] = dataset.select(indices)
            else:
                balanced[name] = dataset

        return balanced

class EmbeddingTrainingPipeline:
    """완전한 임베딩 모델 훈련 파이프라인"""

    def __init__(self, base_model: str = "klue/bert-base"):
        self.base_model = base_model
        self.processor = MultiTaskDataProcessor(base_model)
        self.miner = None

    def load_korean_datasets(self) -> Dict[str, Dataset]:
        """한국어 다중 태스크 데이터셋 로드"""
        datasets = {}

        # 1. KorSTS (Korean STS)
        korsts_data = self._load_korsts()
        if korsts_data:
            datasets['sts'] = self.processor.process_sts_data(korsts_data)

        # 2. KorNLI (Korean NLI) 
        kornli_data = self._load_kornli()
        if kornli_data:
            datasets['nli'] = self.processor.process_nli_data(kornli_data)

        # 3. Korean Question-Answer pairs
        korqa_data = self._load_korqa()
        if korqa_data:
            datasets['qa'] = Dataset.from_list(korqa_data)

        # 4. Korean Classification data
        korclass_data = self._load_korean_classification()
        if korclass_data:
            datasets['classification'] = Dataset.from_list(korclass_data)

        return datasets

    def _load_korsts(self) -> List[Dict]:
        """KorSTS 데이터 로드 (예시)"""
        # 실제로는 Hugging Face datasets나 로컬 파일에서 로드
        return [
            {'sentence1': '오늘 날씨가 좋다', 'sentence2': '날씨가 맑다', 'score': 4.2},
            {'sentence1': '고양이가 귀엽다', 'sentence2': '강아지가 예쁘다', 'score': 2.1},
            # ... 더 많은 데이터
        ]

    def _load_kornli(self) -> List[Dict]:
        """KorNLI 데이터 로드 (예시)"""
        return [
            {'premise': '사람이 걷고 있다', 'hypothesis': '사람이 이동하고 있다', 'label': 'entailment'},
            {'premise': '비가 온다', 'hypothesis': '눈이 온다', 'label': 'contradiction'},
            # ... 더 많은 데이터
        ]

    def _load_korqa(self) -> List[Dict]:
        """한국어 QA 데이터 로드 (예시)"""
        return [
            {'query': '서울의 수도는?', 'answer': '서울은 대한민국의 수도입니다.'},
            {'query': 'AI란 무엇인가?', 'answer': '인공지능(AI)은 컴퓨터가 인간의 지능을 모방하는 기술입니다.'},
            # ... 더 많은 데이터
        ]

    def _load_korean_classification(self) -> List[Dict]:
        """한국어 분류 데이터 로드 (예시)"""
        return [
            {'text': '이 영화는 정말 재미있어요', 'label': 'positive'},
            {'text': '서비스가 별로였습니다', 'label': 'negative'},
            # ... 더 많은 데이터
        ]

    def setup_hard_negative_mining(self, model: SentenceTransformer, passages: List[str]):
        """하드 네거티브 마이닝 설정"""
        self.miner = HardNegativeMiner(model)
        self.miner.build_index(passages)

    def create_training_config(self) -> Dict:
        """훈련 설정 생성"""
        return {
            # 모델 설정
            'base_model': self.base_model,
            'max_seq_length': 512,

            # 훈련 설정
            'epochs': 3,
            'batch_size': 32,
            'learning_rate': 2e-5,
            'warmup_ratio': 0.1,

            # Multi-task 설정
            'task_weights': {
                'sts': 0.3,
                'nli': 0.2, 
                'retrieval': 0.3,
                'classification': 0.2
            },

            # 하드 네거티브 설정
            'use_hard_negatives': True,
            'negative_ratio': 0.3,
            'hard_negative_threshold': (0.3, 0.8),

            # 평가 설정
            'eval_steps': 500,
            'save_steps': 1000,
            'eval_strategy': 'steps'
        }

    def run_full_pipeline(self):
        """완전한 훈련 파이프라인 실행"""
        print("🚀 Starting Multi-task Embedding Training Pipeline")

        # 1. 데이터 로드
        print("📊 Loading datasets...")
        datasets = self.load_korean_datasets()

        # 2. 데이터 균형 맞추기
        print("⚖️ Balancing datasets...")
        balanced_datasets = self.processor.balance_dataset(datasets)

        # 3. 모델 초기화
        print("🤖 Initializing model...")
        model = SentenceTransformer(self.base_model)

        # 4. 하드 네거티브 마이닝 설정
        if 'qa' in balanced_datasets:
            print("🔍 Setting up hard negative mining...")
            passages = [item['answer'] for item in balanced_datasets['qa']]
            self.setup_hard_negative_mining(model, passages)

        # 5. 훈련 설정 생성
        config = self.create_training_config()

        print("✅ Pipeline setup complete!")
        print(f"📈 Training will use {len(balanced_datasets)} tasks:")
        for name, dataset in balanced_datasets.items():
            print(f"   - {name}: {len(dataset)} samples")

        return model, balanced_datasets, config

# 사용 예시
if __name__ == "__main__":
    # 파이프라인 초기화
    pipeline = EmbeddingTrainingPipeline("klue/bert-base")

    # 전체 파이프라인 실행
    model, datasets, config = pipeline.run_full_pipeline()

    # 설정 출력
    print("\n📋 Training Configuration:")
    for key, value in config.items():
        print(f"   {key}: {value}")

Multi-task learning과 advanced techniques(Hard Negatives, Instruction Tuning)을 결합할 때 모든 task에서 균형잡힌 높은 성능을 달성할 수 있습니다.  

 

각 task마다 다른 수렴 패턴을 보이며, NLI는 빠르게 수렴하고 STS는 상대적으로 오래 걸리는 것을 확인할 수 있습니다.

⚙️ 하이퍼파라미터 최적화 가이드

# Multi-task Embedding Model 하이퍼파라미터 최적화 가이드

import optuna
import torch
from typing import Dict, Any
from sentence_transformers import SentenceTransformer
import numpy as np

class EmbeddingHyperparameterOptimizer:
    """임베딩 모델 하이퍼파라미터 최적화"""

    def __init__(self, base_model: str = "klue/bert-base"):
        self.base_model = base_model
        self.best_params = None
        self.best_score = -float('inf')

    def suggest_hyperparameters(self, trial: optuna.Trial) -> Dict[str, Any]:
        """Optuna trial을 사용한 하이퍼파라미터 제안"""

        params = {
            # 학습률 관련
            'learning_rate': trial.suggest_float('learning_rate', 1e-6, 1e-3, log=True),
            'warmup_ratio': trial.suggest_float('warmup_ratio', 0.0, 0.3),
            'lr_scheduler_type': trial.suggest_categorical('lr_scheduler_type', 
                ['linear', 'cosine', 'cosine_with_restarts', 'polynomial']),

            # 배치 크기
            'per_device_train_batch_size': trial.suggest_categorical('batch_size', [16, 32, 64, 128]),

            # 정규화
            'weight_decay': trial.suggest_float('weight_decay', 0.0, 0.1),

            # Temperature parameter for contrastive learning  
            'temperature': trial.suggest_float('temperature', 0.01, 0.1),

            # Multi-task weighting
            'sts_weight': trial.suggest_float('sts_weight', 0.1, 0.5),
            'nli_weight': trial.suggest_float('nli_weight', 0.1, 0.5),
            'retrieval_weight': trial.suggest_float('retrieval_weight', 0.2, 0.6),
            'classification_weight': trial.suggest_float('classification_weight', 0.1, 0.4),

            # Hard negative mining
            'hard_neg_ratio': trial.suggest_float('hard_neg_ratio', 0.2, 0.8),
            'hard_neg_threshold_min': trial.suggest_float('hard_neg_threshold_min', 0.1, 0.4),
            'hard_neg_threshold_max': trial.suggest_float('hard_neg_threshold_max', 0.6, 0.9),
        }

        # 가중치 정규화
        total_weight = (params['sts_weight'] + params['nli_weight'] + 
                       params['retrieval_weight'] + params['classification_weight'])
        params['sts_weight'] /= total_weight
        params['nli_weight'] /= total_weight  
        params['retrieval_weight'] /= total_weight
        params['classification_weight'] /= total_weight

        return params

    def objective(self, trial: optuna.Trial) -> float:
        """최적화 목적 함수"""
        params = self.suggest_hyperparameters(trial)

        # 모델 훈련 (실제로는 여기서 전체 훈련 파이프라인 실행)
        score = self.train_and_evaluate(params)

        return score

    def train_and_evaluate(self, params: Dict[str, Any]) -> float:
        """주어진 파라미터로 모델 훈련 및 평가"""
        # 실제 훈련 코드는 여기에 구현
        # 간단한 모의 점수 반환
        return np.random.random()

    def optimize(self, n_trials: int = 50) -> Dict[str, Any]:
        """하이퍼파라미터 최적화 실행"""
        study = optuna.create_study(direction='maximize')
        study.optimize(self.objective, n_trials=n_trials)

        self.best_params = study.best_params
        self.best_score = study.best_value

        return self.best_params

class TrainingScheduler:
    """훈련 스케쥴 관리"""

    def __init__(self):
        self.schedules = self.create_training_schedules()

    def create_training_schedules(self) -> Dict[str, Dict]:
        """다양한 훈련 스케쥴 정의"""

        schedules = {
            # 빠른 프로토타이핑용
            'quick': {
                'total_epochs': 1,
                'per_device_train_batch_size': 64,
                'learning_rate': 5e-5,
                'warmup_ratio': 0.1,
                'eval_steps': 100,
                'description': '빠른 테스트용 - 1 epoch'
            },

            # 표준 훈련
            'standard': {
                'total_epochs': 3,
                'per_device_train_batch_size': 32,
                'learning_rate': 2e-5,
                'warmup_ratio': 0.1,
                'eval_steps': 500,
                'lr_scheduler_type': 'cosine',
                'description': '표준 훈련 - 3 epochs'
            },

            # 고품질 훈련 (시간 오래 걸림)
            'high_quality': {
                'total_epochs': 5,
                'per_device_train_batch_size': 16,
                'learning_rate': 1e-5,
                'warmup_ratio': 0.2,
                'eval_steps': 200,
                'lr_scheduler_type': 'cosine_with_restarts',
                'weight_decay': 0.01,
                'description': '고품질 훈련 - 5 epochs, 낮은 학습률'
            },

            # 대용량 데이터용
            'large_scale': {
                'total_epochs': 2,
                'per_device_train_batch_size': 128,
                'gradient_accumulation_steps': 4,
                'learning_rate': 3e-5,
                'warmup_ratio': 0.05,
                'eval_steps': 1000,
                'fp16': True,
                'dataloader_num_workers': 8,
                'description': '대용량 데이터용 - 큰 배치 크기'
            }
        }

        return schedules

    def get_schedule(self, schedule_name: str) -> Dict:
        """특정 스케쥴 반환"""
        return self.schedules.get(schedule_name, self.schedules['standard'])

    def print_schedules(self):
        """사용 가능한 스케쥴 출력"""
        print("📅 Available Training Schedules:")
        for name, schedule in self.schedules.items():
            print(f"\n🔧 {name.upper()}:")
            print(f"   📝 {schedule['description']}")
            for key, value in schedule.items():
                if key != 'description':
                    print(f"   {key}: {value}")

class ModelSizeRecommender:
    """모델 크기별 추천 설정"""

    @staticmethod
    def get_recommendations() -> Dict[str, Dict]:
        """모델 크기별 추천 설정"""

        recommendations = {
            'small': {
                'models': ['klue/bert-base', 'monologg/kobert', 'beomi/kcbert-base'],
                'batch_size': 64,
                'learning_rate': 3e-5,
                'training_time': '2-4 hours',
                'memory_requirement': '8GB GPU',
                'recommended_for': '프로토타이핑, 빠른 실험',
                'expected_performance': 'Good baseline performance'
            },

            'medium': {
                'models': ['klue/roberta-large', 'beomi/kcbert-large'],
                'batch_size': 32,
                'learning_rate': 2e-5, 
                'training_time': '6-12 hours',
                'memory_requirement': '16GB GPU',
                'recommended_for': '실용적 응용, 균형잡힌 성능',
                'expected_performance': 'Strong performance on most tasks'
            },

            'large': {
                'models': ['microsoft/DialoGPT-large', 'facebook/dpr-reader-multiset-base'],
                'batch_size': 16,
                'learning_rate': 1e-5,
                'training_time': '12-24 hours', 
                'memory_requirement': '24GB+ GPU',
                'recommended_for': '최고 성능이 필요한 프로덕션',
                'expected_performance': 'State-of-the-art performance'
            }
        }

        return recommendations

    @staticmethod
    def print_recommendations():
        """모델 크기별 추천사항 출력"""
        recommendations = ModelSizeRecommender.get_recommendations()

        print("🤖 Model Size Recommendations:")
        for size, rec in recommendations.items():
            print(f"\n📊 {size.upper()} Models:")
            print(f"   🏷️  Models: {', '.join(rec['models'][:2])}...")
            print(f"   📦 Batch Size: {rec['batch_size']}")
            print(f"   📈 Learning Rate: {rec['learning_rate']}")
            print(f"   ⏰ Training Time: {rec['training_time']}")
            print(f"   💾 Memory: {rec['memory_requirement']}")
            print(f"   🎯 Use Case: {rec['recommended_for']}")
            print(f"   🏆 Performance: {rec['expected_performance']}")

class GPUMemoryEstimator:
    """GPU 메모리 사용량 추정"""

    @staticmethod
    def estimate_memory_usage(model_params: int, batch_size: int, seq_length: int = 512) -> Dict[str, float]:
        """GPU 메모리 사용량 추정 (GB 단위)"""

        # 기본 계산식 (근사치)
        model_memory = model_params * 4 / (1024**3)  # 모델 가중치 (FP32)

        # 순전파 시 activation memory
        activation_memory = batch_size * seq_length * 768 * 4 / (1024**3)  # 768은 hidden_size

        # 역전파를 위한 gradient memory
        gradient_memory = model_memory  # gradient는 모델 크기와 동일

        # 옵티마이저 state (Adam 기준)
        optimizer_memory = model_memory * 2  # Adam은 momentum + variance

        total_memory = model_memory + activation_memory + gradient_memory + optimizer_memory

        return {
            'model_memory_gb': round(model_memory, 2),
            'activation_memory_gb': round(activation_memory, 2), 
            'gradient_memory_gb': round(gradient_memory, 2),
            'optimizer_memory_gb': round(optimizer_memory, 2),
            'total_memory_gb': round(total_memory, 2)
        }

    @staticmethod
    def recommend_batch_size(available_memory_gb: int, model_name: str = "bert-base") -> Dict[str, Any]:
        """사용 가능한 메모리에 따른 배치 크기 추천"""

        model_sizes = {
            'bert-base': 110_000_000,
            'bert-large': 340_000_000,
            'roberta-base': 125_000_000,
            'roberta-large': 355_000_000
        }

        model_params = model_sizes.get(model_name, 110_000_000)

        # 다양한 배치 크기에 대해 메모리 사용량 계산
        batch_recommendations = {}

        for batch_size in [8, 16, 32, 64, 128]:
            memory_usage = GPUMemoryEstimator.estimate_memory_usage(model_params, batch_size)

            if memory_usage['total_memory_gb'] <= available_memory_gb * 0.9:  # 90% 여유
                batch_recommendations[batch_size] = memory_usage

        if batch_recommendations:
            recommended_batch = max(batch_recommendations.keys())
            return {
                'recommended_batch_size': recommended_batch,
                'memory_usage': batch_recommendations[recommended_batch],
                'available_batch_sizes': list(batch_recommendations.keys())
            }
        else:
            return {
                'recommended_batch_size': 4,  # 최소 배치 크기
                'warning': f'Available memory ({available_memory_gb}GB) is too small for efficient training'
            }

# 사용 예시 및 실행
def main():
    """메인 실행 함수"""

    print("🎯 Multi-task Embedding Training Configuration Guide")
    print("="*60)

    # 1. 훈련 스케쥴 출력
    scheduler = TrainingScheduler()
    scheduler.print_schedules()

    print("\n" + "="*60)

    # 2. 모델 크기별 추천사항
    ModelSizeRecommender.print_recommendations()

    print("\n" + "="*60)

    # 3. GPU 메모리 추정
    print("\n💾 GPU Memory Estimation Examples:")

    # BERT-base 예시
    memory_usage = GPUMemoryEstimator.estimate_memory_usage(110_000_000, 32)
    print("\n🔍 BERT-base with batch_size=32:")
    for key, value in memory_usage.items():
        print(f"   {key}: {value}")

    # 배치 크기 추천
    recommendation = GPUMemoryEstimator.recommend_batch_size(16, "bert-base")  # 16GB GPU
    print("\n💡 Batch Size Recommendation for 16GB GPU:")
    for key, value in recommendation.items():
        print(f"   {key}: {value}")

    print("\n" + "="*60)
    print("✅ Configuration guide complete!")
    print("💡 Choose appropriate settings based on your hardware and time constraints.")

if __name__ == "__main__":
    main()

📚 완전한 Step-by-Step 가이드

자세한 단계별 가이드와 FAQ:

Multi-task Embedding Model Training - Complete Step-by-Step Guide
🚀 훈련 프로세스 전체 개요
단계별 훈련 과정:
환경 설정 및 데이터 준비

모델 아키텍처 선택

Multi-task 데이터셋 구성

Loss Function 정의

하드 네거티브 마이닝

훈련 실행

모델 평가 및 검증

모델 배포

📋 Step 1: 환경 설정
필수 라이브러리 설치:

pip install torch transformers sentence-transformers datasets
pip install faiss-cpu  # or faiss-gpu for GPU
pip install optuna wandb  # for hyperparameter tuning and logging

GPU 메모리 체크:

import torch
print(f"GPU available: {torch.cuda.is_available()}")
print(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

📊 Step 2: 데이터셋 선택 및 구성
각 Task별 추천 데이터셋:
🎯 STS (Semantic Textual Similarity)

데이터: STS-B, KorSTS, SICK-R

형식: (sentence1, sentence2, similarity_score)

크기: 10K-100K pairs

Loss: CoSENTLoss, MSELoss

🧠 NLI (Natural Language Inference)

데이터: SNLI, MultiNLI, KorNLI, XNLI

형식: (premise, hypothesis, label)

크기: 500K+ examples

Loss: SoftmaxLoss (3-class)

🔍 Retrieval

데이터: MS-MARCO, Natural Questions, HotpotQA

형식: (query, positive_doc, negative_doc)

크기: 100K-1M triplets

Loss: MultipleNegativesRankingLoss (InfoNCE)

📝 Classification

데이터: AG News, IMDb, Amazon Reviews, 한국어 리뷰 데이터

형식: (text, class_label)

크기: 50K-500K examples

Loss: SoftmaxLoss

🗂️ Clustering

데이터: 20 Newsgroups, Reddit posts, ArXiv papers

형식: (text, cluster_id)

크기: 10K-100K examples

Loss: InfoNCE with cluster-based sampling

🏗️ Step 3: 모델 아키텍처 선택
한국어 모델 추천:
초급자/프로토타이핑:

klue/bert-base: 110M params, 8GB GPU

monologg/kobert: 92M params, 경량화

중급자/실용성:

klue/roberta-large: 355M params, 16GB GPU

beomi/kcbert-large: 340M params

고급자/최고성능:

microsoft/DialoGPT-large: 762M params, 24GB+ GPU

Custom model with specialized architecture

⚙️ Step 4: 훈련 설정 구성
기본 하이퍼파라미터:

training_args = {
    'learning_rate': 2e-5,
    'batch_size': 32, 
    'epochs': 3,
    'warmup_ratio': 0.1,
    'weight_decay': 0.01,
    'lr_scheduler_type': 'cosine',
    'fp16': True,  # 메모리 절약
}

Task별 가중치 (중요도에 따라 조절):

task_weights = {
    'retrieval': 0.4,  # 검색 성능이 가장 중요한 경우
    'sts': 0.3,
    'nli': 0.2, 
    'classification': 0.1
}

🔍 Step 5: 하드 네거티브 마이닝
마이닝 전략:
Random Negatives: 기본 베이스라인

BM25 Negatives: 키워드 기반 유사성

Dense Negatives: 현재 모델로 유사한 문서 검색

Cross-encoder Negatives: 더 정확한 teacher 모델 사용

마이닝 비율:
Easy negatives: 30%

Hard negatives: 50%

Very hard negatives: 20%

🚀 Step 6: 실제 훈련 실행
단계별 훈련:
Phase 1: Contrastive Pre-training (옵션)

# 대규모 약지도 데이터로 기본 표현 학습
# 1B+ text pairs, InfoNCE loss
# 수만 steps, 큰 배치 크기 (512-2048)

Phase 2: Multi-task Fine-tuning

# 라벨이 있는 다양한 task 데이터로 특화 학습
# Task-specific losses 조합
# 수천 steps, 적당한 배치 크기 (32-128)

모니터링 지표:
STS: Spearman correlation

NLI: Accuracy

Retrieval: MRR@10, Recall@100

Classification: F1-score

Overall: Weighted average of all metrics

📊 Step 7: 모델 평가
평가 방법:
Task-specific metrics: 각 태스크별 표준 평가 지표

MTEB benchmark: 종합 평가 벤치마크

Downstream evaluation: 실제 사용 사례에서의 성능

Ablation study: 각 컴포넌트별 기여도 분석

평가 데이터셋:
훈련에 사용되지 않은 held-out test sets

도메인별 specialized test sets

Cross-lingual evaluation (다국어 모델인 경우)

🚀 Step 8: 모델 배포
최적화:

# 모델 경량화
model.half()  # FP16 변환
model = torch.jit.script(model)  # TorchScript 변환

# ONNX 변환 (추론 최적화)
torch.onnx.export(model, input_tensor, "model.onnx")

서빙:
Hugging Face Hub: 모델 공유 및 배포

FastAPI/Flask: REST API 서버

Docker: 컨테이너화 배포

Kubernetes: 스케일링 가능한 배포

❓ FAQ - 자주 묻는 질문들
Q1: 얼마나 많은 데이터가 필요한가요?
A: 태스크별로 다르지만 최소 권장량:

STS: 10K+ pairs

NLI: 100K+ examples

Retrieval: 50K+ query-doc pairs

Classification: 10K+ per class

전체적으로 500K+ examples 권장

Q2: 훈련 시간은 얼마나 걸리나요?
A:

Quick (1 epoch): 2-4시간 (BERT-base, 16GB GPU)

Standard (3 epochs): 6-12시간

High-quality (5 epochs): 12-24시간

Large-scale models: 1-3일

Q3: GPU 메모리가 부족할 때 해결 방법은?
A:

배치 크기 줄이기 (32 → 16 → 8)

Gradient accumulation 사용

FP16 mixed precision training

Gradient checkpointing 활성화

더 작은 모델 사용

Q4: 한국어 데이터가 부족할 때는?
A:

Data Augmentation: Back-translation, paraphrasing

Transfer Learning: 영어 모델에서 시작

Synthetic Data: LLM으로 생성

Cross-lingual Training: 다국어 데이터 혼합 사용

Domain Adaptation: 일반 도메인 → 특정 도메인

Q5: 모델 성능이 좋지 않을 때 디버깅 방법은?
A:

데이터 품질 체크: 라벨 오류, 중복 데이터

Loss 곡선 분석: Overfitting, underfitting 확인

Learning rate 조정: 너무 높으면 불안정, 너무 낮으면 느림

Task 간 간섭 확인: 상충하는 objective들

Negative sampling 개선: 하드 네거티브 비율 조정

Q6: Production 환경에서 주의사항은?
A:

Latency 최적화: 모델 크기 vs 성능 트레이드오프

Scalability: 동시 요청 처리 능력

Model versioning: A/B 테스트 가능한 구조

Monitoring: 성능 저하 감지 시스템

Security: 악의적 입력에 대한 방어

Q7: 어떤 평가 지표를 사용해야 하나요?
A:

STS: Spearman correlation (연속적 유사도)

Retrieval: MRR@10, Recall@100, nDCG

Classification: F1-score (불균형 데이터), Accuracy

Clustering: Adjusted Rand Index, V-measure

종합: MTEB benchmark 점수

Q8: 멀티태스크 학습의 장단점은?
A:
장점:

데이터 효율성 향상

일반화 성능 개선

리소스 절약 (하나의 모델)

Task 간 지식 전이

단점:

Task 간 간섭 가능성

하이퍼파라미터 튜닝 복잡

개별 task 최적화 어려움

디버깅 복잡성

Q9: 최신 트렌드는 무엇인가요?
A:

Instruction-following: E5, BGE 스타일 instruction tuning

Mixture of Experts: 태스크별 전문가 모델

Retrieval-Augmented: RAG와 통합

Multimodal: 텍스트 + 이미지 결합

Efficiency: Distillation, pruning, quantization

🎯 최종 추천사항
초보자를 위한 시작점:
klue/bert-base 모델 사용

Standard 훈련 스케쥴 적용

기본 데이터셋들로 시작 (STS + NLI + 하나의 도메인)

간단한 평가 지표로 성능 확인

고급 사용자를 위한 팁:
Custom architecture 실험

Advanced negative mining 전략

Curriculum learning 적용

Multi-stage training pipeline

성공의 핵심:

충분한 고품질 데이터

적절한 하이퍼파라미터 튜닝

체계적인 평가 및 분석

지속적인 모니터링 및 개선

💡 Remember: 완벽한 모델을 한 번에 만들려 하지 말고,
점진적으로 개선해 나가는 것이 중요합니다!

🎮 실습 데모 코드

실제로 multi-task embedding이 어떻게 작동하는지 체험해볼 수 있는 데모:

#!/usr/bin/env python3
"""
Multi-task Embedding Model Training Demo
간단하게 실행해볼 수 있는 데모 코드입니다.
"""

import torch
import numpy as np
from transformers import AutoTokenizer, AutoModel
import torch.nn.functional as F
from typing import List, Dict, Any
import warnings
warnings.filterwarnings("ignore")

class SimpleMultiTaskEmbedding:
    """간단한 Multi-task Embedding 모델 데모"""

    def __init__(self, model_name: str = "klue/bert-base"):
        """모델 초기화"""
        print(f"🤖 Loading model: {model_name}")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.to(self.device)
        print(f"📱 Device: {self.device}")

    def encode(self, texts: List[str], batch_size: int = 32) -> torch.Tensor:
        """텍스트를 임베딩으로 변환"""
        self.model.eval()
        embeddings = []

        with torch.no_grad():
            for i in range(0, len(texts), batch_size):
                batch_texts = texts[i:i + batch_size]

                # 토크나이징
                inputs = self.tokenizer(
                    batch_texts,
                    padding=True,
                    truncation=True,
                    max_length=512,
                    return_tensors="pt"
                ).to(self.device)

                # 임베딩 생성
                outputs = self.model(**inputs)
                # Mean pooling
                batch_embeddings = self.mean_pooling(outputs, inputs['attention_mask'])
                embeddings.append(batch_embeddings)

        return torch.cat(embeddings, dim=0)

    def mean_pooling(self, model_output, attention_mask):
        """Mean pooling 적용"""
        token_embeddings = model_output[0]
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

    def compute_similarity(self, embeddings1: torch.Tensor, embeddings2: torch.Tensor) -> torch.Tensor:
        """코사인 유사도 계산"""
        return F.cosine_similarity(embeddings1, embeddings2, dim=1)

def demo_sts_task():
    """STS (Semantic Textual Similarity) 태스크 데모"""
    print("\n🎯 STS (Semantic Textual Similarity) Demo")
    print("-" * 50)

    model = SimpleMultiTaskEmbedding()

    # 예시 문장 쌍들
    sentence_pairs = [
        ("오늘 날씨가 좋다", "날씨가 맑다"),
        ("고양이가 귀엽다", "강아지가 예쁘다"),
        ("인공지능을 공부한다", "AI를 학습한다"),
        ("사과를 먹었다", "자동차를 운전했다"),
        ("안녕하세요", "안녕히 가세요")
    ]

    print("📊 문장 유사도 분석 결과:")
    for sent1, sent2 in sentence_pairs:
        embeddings1 = model.encode([sent1])
        embeddings2 = model.encode([sent2])
        similarity = model.compute_similarity(embeddings1, embeddings2).item()

        print(f"'{sent1}' vs '{sent2}'")
        print(f"   유사도: {similarity:.4f} {'😊' if similarity > 0.7 else '🤔' if similarity > 0.4 else '😐'}")
        print()

def demo_retrieval_task():
    """Retrieval 태스크 데모"""
    print("\n🔍 Retrieval Demo")
    print("-" * 50)

    model = SimpleMultiTaskEmbedding()

    # 문서 컬렉션
    documents = [
        "파이썬은 프로그래밍 언어입니다",
        "서울은 대한민국의 수도입니다", 
        "인공지능은 미래 기술입니다",
        "김치는 한국의 전통 음식입니다",
        "축구는 전 세계에서 인기 있는 스포츠입니다"
    ]

    # 쿼리
    queries = [
        "프로그래밍에 대해 알려주세요",
        "한국의 수도는 어디인가요?",
        "AI 기술이 궁금합니다"
    ]

    # 문서 임베딩 생성
    print("📚 문서 임베딩 생성 중...")
    doc_embeddings = model.encode(documents)

    print("\n🔎 검색 결과:")
    for query in queries:
        query_embedding = model.encode([query])

        # 모든 문서와의 유사도 계산
        similarities = model.compute_similarity(
            query_embedding.repeat(len(documents), 1), 
            doc_embeddings
        )

        # 상위 2개 문서 선택
        top_indices = torch.argsort(similarities, descending=True)[:2]

        print(f"\n쿼리: '{query}'")
        for i, idx in enumerate(top_indices):
            print(f"  {i+1}. {documents[idx]} (유사도: {similarities[idx]:.4f})")

def demo_classification_task():
    """Classification 태스크 데모 (감정 분석)"""
    print("\n📝 Classification Demo - 감정 분석")
    print("-" * 50)

    model = SimpleMultiTaskEmbedding()

    # 예시 텍스트들
    texts = [
        "이 영화는 정말 재미있어요!",
        "서비스가 너무 별로였습니다",
        "그냥 보통이에요",
        "최고의 경험이었습니다",
        "다시는 이용하지 않을 것 같아요"
    ]

    # 감정 레이블 (실제로는 학습 데이터에서 학습)
    positive_template = "이것은 긍정적인 감정입니다"
    negative_template = "이것은 부정적인 감정입니다"
    neutral_template = "이것은 중립적인 감정입니다"

    templates = [positive_template, negative_template, neutral_template]
    labels = ["긍정", "부정", "중립"]

    print("🎭 감정 분석 결과:")

    # 템플릿 임베딩 생성
    template_embeddings = model.encode(templates)

    for text in texts:
        text_embedding = model.encode([text])

        # 각 템플릿과의 유사도 계산
        similarities = model.compute_similarity(
            text_embedding.repeat(len(templates), 1),
            template_embeddings
        )

        # 가장 유사한 템플릿 선택
        predicted_idx = torch.argmax(similarities).item()
        confidence = similarities[predicted_idx].item()

        emoji = "😊" if predicted_idx == 0 else "😠" if predicted_idx == 1 else "😐"
        print(f"'{text}'")
        print(f"   예측: {labels[predicted_idx]} {emoji} (확신도: {confidence:.4f})")
        print()

def demo_clustering_task():
    """Clustering 태스크 데모"""
    print("\n🗂️ Clustering Demo")
    print("-" * 50)

    model = SimpleMultiTaskEmbedding()

    # 서로 다른 주제의 텍스트들
    texts = [
        # 스포츠 관련
        "축구는 재미있는 스포츠입니다",
        "야구 경기를 보러 갔습니다", 
        "올림픽 선수들이 훌륭합니다",

        # 음식 관련
        "김치찌개가 맛있어요",
        "파스타를 만들어 먹었습니다",
        "디저트로 케이크를 먹었어요",

        # 기술 관련
        "인공지능이 발전하고 있습니다",
        "새로운 스마트폰이 출시되었어요",
        "프로그래밍을 배우고 있습니다"
    ]

    print("📊 텍스트 클러스터링 결과:")

    # 임베딩 생성
    embeddings = model.encode(texts)

    # 간단한 클러스터링 (유사도 기반)
    clusters = {}
    cluster_id = 0
    used_indices = set()

    for i, embedding1 in enumerate(embeddings):
        if i in used_indices:
            continue

        cluster = [i]
        used_indices.add(i)

        # 유사한 텍스트들 찾기
        for j, embedding2 in enumerate(embeddings):
            if j != i and j not in used_indices:
                similarity = F.cosine_similarity(
                    embedding1.unsqueeze(0), 
                    embedding2.unsqueeze(0)
                ).item()

                if similarity > 0.5:  # 임계값
                    cluster.append(j)
                    used_indices.add(j)

        clusters[f"클러스터 {cluster_id + 1}"] = cluster
        cluster_id += 1

    for cluster_name, indices in clusters.items():
        print(f"\n{cluster_name}:")
        for idx in indices:
            print(f"  - {texts[idx]}")

def main():
    """메인 데모 실행 함수"""
    print("🚀 Multi-task Embedding Model Demo")
    print("=" * 60)
    print("다양한 NLP 태스크에서 임베딩 모델을 어떻게 사용하는지 보여드립니다!")

    try:
        # 각 태스크별 데모 실행
        demo_sts_task()
        demo_retrieval_task() 
        demo_classification_task()
        demo_clustering_task()

        print("\n" + "=" * 60)
        print("✅ 모든 데모가 완료되었습니다!")
        print("💡 실제 Multi-task 학습을 통해 이 모든 태스크에서 더 좋은 성능을 얻을 수 있습니다.")

    except Exception as e:
        print(f"❌ 오류 발생: {e}")
        print("💡 인터넷 연결과 필요한 라이브러리가 설치되어 있는지 확인하세요.")
        print("   pip install torch transformers")

if __name__ == "__main__":
    main()

🔑 핵심 훈련 전략

2단계 훈련 접근법

1.       Stage 1: Contrastive Pre-training (선택사항)

o    대규모 약지도 데이터 (1B+ text pairs)

o    InfoNCE Loss 사용

o    기본적인 텍스트 표현 학습

2.      Stage 2: Multi-task Fine-tuning

o    다양한 labeled 데이터 조합

o    Task-specific losses 결합

o    특화된 능력 개발

Hard Negative Mining

# 하드 네거티브 마이닝 전략
hard_negative_strategy = {
    'easy_negatives': 0.3,    # 30%
    'hard_negatives': 0.5,    # 50%  
    'very_hard_negatives': 0.2 # 20%
}

추천 하이퍼파라미터

recommended_config = {
    'learning_rate': 2e-5,
    'batch_size': 32,
    'epochs': 3,
    'warmup_ratio': 0.1,
    'weight_decay': 0.01,
    'temperature': 0.05,  # InfoNCE용
    
    # Task 가중치
    'task_weights': {
        'retrieval': 0.4,
        'sts': 0.3,
        'nli': 0.2,
        'classification': 0.1
    }
}

💡 성공을 위한 핵심 팁

  1. 데이터 품질이 최우선: 고품질 데이터가 모델 성능을 좌우합니다
  2. 점진적 개선: 완벽한 모델을 한번에 만들려 하지 말고 단계적으로 개선
  3. 체계적 평가: MTEB benchmark 등을 사용한 객관적 평가
  4. 하드웨어 고려: GPU 메모리에 맞는 배치 크기와 모델 선택
  5. 지속적 모니터링: 훈련 중 각 task별 성능 추적

이 가이드를 따라 진행하시면 모든 NLP task에서 우수한 성능을 보이는 universal embedding 모델을 만들 수 있습니다. 궁금한 점이 있으시면 언제든 질문해 주세요!

 

728x90
저작자표시 비영리 (새창열림)

'인공지능 > 자연어 처리' 카테고리의 다른 글

허깅페이스 3 기초 - Audio Feature Extractors  (0) 2025.11.15
허깅페이스 1 기초 - 모델 부르기, 모델 공유하기, 모델 구성 요소 변경하기  (0) 2025.11.12
LLM Pruning to Encoder - Large Language Models Are Overparameterized Text Encoders  (2) 2025.07.17
Generation + Embedding 세미나 준비  (0) 2025.06.12
LLM 평가 코드 분석하기 - lm evaluation harness  (0) 2025.05.14

'인공지능/자연어 처리'의 다른글

  • 현재글Embedding 모델 학습하기 - Sentence Transformer Trainer

관련글

티스토리툴바