LLM Research2026.10.3010 min read

DPO Loss 직접 구현 + TRL DPOTrainer: 수식을 코드로 옮기고 라이브러리 결과와 대조하기

DPO(Direct Preference Optimization)는 "좋은 답변"과 "나쁜 답변" 쌍만으로 모델을 선호 방향으로 옮기는 학습 방법이다. 오늘은 DPO loss를 수식에서 PyTorch 코드로 직접 구현하고, 숫자 몇 개로 gradient 방향이 맞는지 확인한다. 이어서 TRL의 DPOTrainer로 같은 데이터를 학습시켜 내 구현과 결과를 비

By 송민성

오늘의 목표

DPO loss를 논문 수식에서 PyTorch 코드로 직접 옮겨 구현하고, 아주 작은 숫자 예제로 gradient(기울기, 손실을 키우거나 줄이는 방향을 알려주는 값)의 방향이 상식과 맞는지 확인한다. 그다음 Hugging Face TRL 라이브러리의 DPOTrainer로 같은 데이터를 학습시켜, 내 손실값·학습 결과와 라이브러리 결과가 얼마나 일치하는지 대조한다.

개념과 직관

선호 데이터(preference data)란

지난 Day들에서 다룬 선호 데이터를 떠올려 보자. 한 질문(prompt)에 대해 답변 두 개가 있고, 사람이(또는 다른 모델이) "이쪽이 더 낫다"고 표시한 데이터다.

prompt: 질문
chosen: 더 좋다고 고른 답변 (y_w, winner)
rejected: 덜 좋다고 고른 답변 (y_l, loser)

DPO의 직관

Reward Model 방식은 "점수 매기는 모델"을 따로 학습한 뒤 그 점수로 다시 언어 모델을 학습하는 2단계다. DPO는 이 중간 단계를 건너뛰고, 언어 모델 하나를 선호 쌍으로 직접 학습한다.

비유를 들자. 요리 학원생(학습 중인 모델, policy)이 있고, 학원 입학 때의 실력을 기록해 둔 사본(reference 모델, 얼려 둔 원본)이 있다. 심사위원이 "A 요리가 B 요리보다 낫다"고 알려 준다. 목표는 이렇다.

원래 실력(reference)과 비교해서 A를 만들 확률은 더 높이고
B를 만들 확률은 더 낮추되
원래 실력에서 너무 멀리 도망가지는 않는다.

"reference 대비 얼마나 변했나"가 핵심이다. 절대 확률이 아니라 변화량(비율)을 본다.

최소 수식

로그 확률(log probability, 모델이 어떤 답변을 낼 확률에 log를 취한 값. 문장 전체는 토큰별 log prob의 합)을 쓴다.

text

policy_logratio    = logp_policy(chosen)    - logp_policy(rejected)
reference_logratio = logp_ref(chosen)       - logp_ref(rejected)

logits = beta * (policy_logratio - reference_logratio)

loss = -log(sigmoid(logits))

하나씩 풀면:

policy_logratio: 지금 모델이 chosen을 rejected보다 얼마나 더 선호하는가
reference_logratio: 원본 모델은 원래 얼마나 선호했는가
둘의 차이: 학습으로 "선호가 얼마나 좋은 방향으로 이동했는가"
sigmoid(시그모이드, 어떤 실수든 0~1 사이로 눌러 주는 함수): 1 / (1 + exp(-x))
-log(sigmoid(x)): x가 클수록(이동이 클수록) loss가 0에 가까워지고, x가 음수면 loss가 커진다.

beta의 효과

beta는 "reference에서 얼마나 멀어지는 걸 허용하나"를 조절하는 숫자다.

beta가 크면: 작은 변화에도 logits가 크게 반응한다. 같은 이동량이어도 loss가 빨리 줄어들기 때문에 모델이 reference 근처에 머무르려는 경향이 강해진다(보수적).
beta가 작으면: logits 반응이 둔하다. loss를 줄이려면 더 크게 움직여야 하므로 reference에서 멀리 벗어나기 쉽다.

일반적으로 0.1 근처 값이 자주 쓰이지만, 정답은 데이터·모델마다 다르니 오늘 실험에서 직접 비교해 본다.

reference 고정과 log prob 캐싱

reference 모델은 학습하지 않는다. 파라미터를 얼려 둔다(requires_grad=False, gradient 계산 대상에서 제외). reference도 같이 움직이면 "기준"이 아니게 된다.
reference의 log prob은 학습 내내 같은 입력에 같은 값이다. 그래서 미리 한 번 계산해서 저장(캐싱, caching) 해 두면 매 스텝 reference forward 계산과 메모리를 아낄 수 있다.

코드로 직접 해보기

1단계: DPO loss 구현

python

import torch
import torch.nn.functional as F

def dpo_loss(policy_chosen_logp, policy_rejected_logp,
             ref_chosen_logp, ref_rejected_logp, beta=0.1):
    policy_logratio = policy_chosen_logp - policy_rejected_logp
    ref_logratio = ref_chosen_logp - ref_rejected_logp
    logits = beta * (policy_logratio - ref_logratio)
    loss = -F.logsigmoid(logits)   # -log(sigmoid(x))를 수치적으로 안전하게
    return loss.mean(), logits

torch.log(torch.sigmoid(x)) 대신 F.logsigmoid를 쓰는 이유: x가 아주 작은 음수일 때 sigmoid가 0에 붙어 log가 -inf가 되는 문제를 피한다.

2단계: 유한 예제로 gradient 방향 확인

모델 없이 숫자 몇 개로 먼저 검증한다. policy log prob을 "학습 가능한 값"으로 두고 gradient를 본다.

python

policy_chosen = torch.tensor([-10.0], requires_grad=True)
policy_rejected = torch.tensor([-12.0], requires_grad=True)
ref_chosen = torch.tensor([-10.0])
ref_rejected = torch.tensor([-12.0])

loss, logits = dpo_loss(policy_chosen, policy_rejected,
                        ref_chosen, ref_rejected, beta=0.1)
loss.backward()

print("loss:", loss.item())              # 처음엔 logits=0 이므로 log(2) ≈ 0.6931
print("grad chosen:", policy_chosen.grad.item())
print("grad rejected:", policy_rejected.grad.item())

기대하는 결과:

loss ≈ 0.6931 (policy = reference이면 logits = 0, -log(0.5) = log 2)
grad chosen은 음수, grad rejected는 양수

해석: 경사하강은 gradient의 반대 방향으로 값을 옮긴다. chosen의 gradient가 음수이므로 chosen log prob은 올라가고, rejected는 양수이므로 내려간다. 우리가 원한 방향이다. 이 부호가 반대로 나오면 chosen/rejected를 뒤바꿨거나 부호 실수가 있다는 신호다.

추가로, chosen을 -8로 올려 보고 loss가 줄어드는지도 확인한다.

python

for c in [-10.0, -9.0, -8.0]:
    l, _ = dpo_loss(torch.tensor([c]), torch.tensor([-12.0]),
                    ref_chosen, ref_rejected, beta=0.1)
    print(c, l.item())   # c가 커질수록 loss가 감소해야 함

3단계: 실제 모델에서 log prob 계산

문장 전체의 log prob은 "답변 토큰"의 log prob 합이다. prompt 부분은 빼야 한다.

python

def sequence_logp(model, input_ids, attention_mask, prompt_len):
    logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
    # 위치 t의 logits는 t+1번째 토큰을 예측하므로 한 칸 밀어서 맞춘다
    logits = logits[:, :-1, :]
    labels = input_ids[:, 1:]
    logp = torch.log_softmax(logits, dim=-1)
    token_logp = logp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)

    mask = attention_mask[:, 1:].clone()
    mask[:, :prompt_len - 1] = 0     # prompt 구간 제외 (배치 내 prompt 길이가 같다고 가정)
    return (token_logp * mask).sum(dim=-1)

reference는 torch.no_grad() 안에서 한 번만 계산해 저장한다.

python

with torch.no_grad():
    ref_chosen_logp = sequence_logp(ref_model, chosen_ids, chosen_mask, prompt_len)
    ref_rejected_logp = sequence_logp(ref_model, rejected_ids, rejected_mask, prompt_len)

작은 모델(예: 수억 파라미터 이하 instruct/base 모델)과 수십~수백 개의 선호 쌍이면 노트북 환경에서도 흐름을 확인할 수 있다. 모델·데이터 이름은 본인 환경에서 받을 수 있는 것으로 고른다.

4단계: TRL DPOTrainer로 같은 데이터 재현

TRL은 버전에 따라 인자 이름이 바뀐다(예: tokenizer → processing_class, DPOConfig의 필드 등). 아래는 큰 흐름이고, 세부는 설치한 버전의 공식 문서로 확인한다.

python

from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import DPOConfig, DPOTrainer

model_name = "여기에_작은_모델_이름"   # 직접 선택
tok = AutoTokenizer.from_pretrained(model_name)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(model_name)

train_ds = Dataset.from_list([
    {"prompt": "...", "chosen": "...", "rejected": "..."},
    # 내 구현에 쓴 것과 같은 데이터
])

args = DPOConfig(output_dir="dpo_out", beta=0.1,
                 per_device_train_batch_size=2,
                 num_train_epochs=1, learning_rate=5e-7,
                 logging_steps=1, report_to="none")

trainer = DPOTrainer(model=model, args=args,
                     train_dataset=train_ds, processing_class=tok)
trainer.train()

ref_model을 따로 안 넘기면 TRL이 내부에서 초기 모델의 복사본을 reference로 쓴다(버전에 따라 동작이 다를 수 있으니 문서 확인).

오늘의 실험

실험 1: 내 loss vs TRL loss 대조 (학습 전)

같은 모델, 같은 데이터 1배치에 대해 내 dpo_loss로 계산한 값과, TRL 첫 스텝의 loss 로그를 비교한다.
학습 전에는 policy = reference이므로 두 값 모두 약 0.693이어야 한다. 여기서 어긋나면 log prob 계산(토큰 시프트, prompt 마스킹)부터 의심한다.

실험 2: beta 바꿔 보기

beta를 0.05, 0.1, 0.5로 바꿔 같은 스텝 수만큼 학습한다.
관찰: loss 곡선, 그리고 학습 후 policy_logratio - reference_logratio의 크기(reference에서 얼마나 멀어졌는가).

실험 3: 결과 대조

학습 후 같은 prompt 몇 개에 대해 reference와 학습된 모델의 답변을 나란히 생성해 눈으로 비교한다.
chosen log prob이 올랐는지, rejected log prob이 내려갔는지 숫자로 확인한다.

기록 표 (직접 채운다)

구현	beta	첫 스텝 loss	마지막 loss	chosen logp 변화	rejected logp 변화	메모
내 구현	0.1
TRL	0.1
TRL	0.05
TRL	0.5

완전히 같은 숫자가 나오지 않을 수 있다. 배치 구성, 길이 처리, 패딩, 랜덤 시드, 정밀도(fp16/bf16) 차이 때문이다. "왜 다른가"를 한 줄이라도 적어 두는 것이 이 실험의 핵심이다.

왜 이걸 배우나 (LLM / Post-training 연결)

DPO는 SFT(지도 미세조정) 다음 단계에서 모델을 사람의 선호에 맞추는 대표적인 post-training 기법이고, 이후 Project 3에서 직접 만든 선호 데이터로 학습·평가하는 데 그대로 쓰인다. LLM Post-training 업무에서는 "라이브러리가 돌아간다"에서 끝나지 않고 loss가 어떻게 계산되는지, beta·reference·데이터 품질이 결과에 어떻게 영향을 주는지 설명하고 디버깅할 수 있어야 하므로, 오늘처럼 직접 구현해서 대조해 본 경험이 바탕이 된다.

초보자가 막히는 지점

chosen/rejected 뒤바뀜: gradient 부호가 반대로 나오거나 loss가 오히려 증가한다. 2단계의 부호 확인이 이 실수를 잡아 준다.
토큰 한 칸 시프트 누락: 모델 출력 logits[t]는 t+1번째 토큰의 예측이다. 시프트를 빼먹으면 log prob이 엉뚱한 값이 되고 TRL과 값이 안 맞는다.
prompt 토큰까지 합산: DPO는 보통 답변 부분의 log prob만 쓴다. prompt를 포함하면 chosen/rejected에 공통인 부분이 섞인다. 마스크를 점검한다.
reference가 같이 학습됨: ref_model을 no_grad/eval로 두지 않으면 기준이 움직인다. 첫 스텝 loss가 0.693인지 확인한다.
learning rate 과다: DPO는 SFT보다 훨씬 작은 learning rate(예: 1e-6 ~ 5e-7 수준)를 쓰는 경우가 많다. 크면 출력이 깨지거나 loss가 불안정해진다.
loss가 내려가도 좋은 모델은 아니다: 선호 쌍을 외우는 과적합일 수 있다. 숫자만 보지 말고 생성 결과도 직접 읽는다.
TRL 버전 차이: 인자 이름 에러(unexpected keyword argument)가 나면 설치된 버전의 문서와 help(DPOConfig)를 확인한다.
GPU 메모리 부족: 모델 2개(policy, reference)가 올라가므로 SFT보다 메모리를 더 쓴다. 더 작은 모델, 짧은 max_length, 작은 배치로 줄인다.

알고리즘 (병행 1시간)

LeetCode Medium 2문제(DP, 동적 계획법: 큰 문제를 작은 문제로 나눠 답을 저장해 두고 재사용하는 방법)를 푼다. 문제는 아래 기준으로 직접 고른다.

1차원 DP 1문제 (예: 계단 오르기·집 도둑질 계열처럼 dp[i]가 앞 몇 칸에만 의존하는 유형)
2차원 DP 또는 문자열 DP 1문제 (예: 격자 경로, 부분수열 계열)

풀이마다 아래를 기록한다.

접근: dp[i](또는 dp[i][j])가 무엇을 뜻하는지 한 문장, 점화식, 초기값
시간복잡도 / 공간복잡도: 그리고 공간을 줄일 수 있는지
틀린 이유: 인덱스 범위, 초기값 누락, 점화식 오류 등 틀렸다면 구체적으로

참고: DP의 "작은 문제 결과를 저장해 재사용"하는 감각은 오늘의 log prob 캐싱과 닮았다.