KTO — 선호 없이 이분 점수로 정렬

KTO

선호 쌍 데이터는 모으기 어려워요. **KTO(Kahneman-Tversky Optimization)**는 '이 답은 괜찮다/아니다'라는 단순한 이분 신호만으로 정렬하는 방법이에요. 행동경제학의 전망 이론(prospect theory)에서 아이디어를 빌려, 인간이 손실에 민감하다는 점을 반영한 손실을 써요.

왜 유용한가

  • 실전에서 선호 쌍(쌍을 이룬 좋은/나쁜 답) 데이터는 값비싸요.
  • KTO는 각 응답이 좋은지 나쁜지만 알면 되므로 데이터 수집이 훨씬 쉬워요.
  • 1B~30B 규모에서 선호 기반 방법과 맞먹거나 더 나은 성능을 내기도 해요.

퀵스타트

from datasets import load_dataset
from trl.experimental.kto import KTOConfig, KTOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
train_dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

training_args = KTOConfig(output_dir="Qwen2-0.5B-KTO")
trainer = KTOTrainer(model=model, args=training_args, processing_class=tokenizer, train_dataset=train_dataset)
trainer.train()

데이터 형식

KTO는 unpaired 선호 데이터가 필요해요. label=True(좋음)/label=False(나쁨)로 표시된 응답을 쓰고, 쌍 데이터가 들어오면 자동으로 분리해서 변환해요.

  • 나쁜(rejected) 데이터만 쓸 땐 보수적인 학습률을 권장해요.
  • 보상 마진이 증가하면 모델이 더 좋은 응답을 만들고 있다는 신호예요.

더 알아보기