HOME » DPO 톺아보기
DPO 톺아보기
개요
2023년 5월에 공개된 DPO 논문1은 그간 PPO 중심의 RLHF 학습 방식이 가진 번거로움을 개선하여 큰 반향을 일으켰습니다. 현재는 DPO와 함께 APO 같은 변형들이 널리 활용되면서 벌써 고전 알고리즘처럼 보일 정도입니다. 하지만 사실 DPO는 등장한 지 이제 고작 2년 남짓밖에 되지 않았죠. 그렇다면 이렇게 큰 반향을 일으킨 DPO 알고리즘의 상세한 내용을 다시 한 번 살펴보겠습니다.
내용
가장 먼저 눈에 띄는 부분은 어떻게 훨씬 더 단순한 DPO가 RLHF와 동일한 성능을 낼 수 있느냐는 점입니다.
수식
먼저 논문 저자들은 오래된 수식을 하나 꺼내듭니다. 인간의 선호도를 모델링하는 수식으로, 1950년대에 소개된 브래들리-테리 모델(1952)2이죠. \(y_1\), \(y_2\) 두 개의 응답 중 \(y_1\)의 리워드가 더 높을 확률을 브래들리-테리 확률모델로 정의하면 다음과 같습니다.
\[p^\ast (y_1 \succ y_2 \vert x) = \frac{\exp (r^\ast (x, y_1))}{\exp (r^\ast (x, y_1)) + \exp (r^\ast (x, y_2))}\]\(p^\ast (y_1 \succ y_2 \vert x)\)가 \(p^\ast (y_2 \succ y_1 \vert x)\)보다 높은 확률을 갖도록 NLLLossNegative Log Likelihood Loss를 다음과 같이 정의할 수 있습니다.
\[\mathcal{L}_R (r_\phi, \mathcal{D}) = - \mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} [\log \sigma (r_\phi (x, y_w) - r_\phi (x, y_l))]\]리워드 모델 \(r_\phi\)는 휴먼 피드백을 받아 점진적으로 업데이트됩니다.
RL 단계에서는 학습 안정성을 높이기 위해 기존 모델을 레퍼런스 \(\pi_\textrm{ref}\)로 두고, 학습 중인 모델 \(\pi_\theta\)와의 편차를 다음과 같이 \(\beta\)로 제어하면서 학습합니다.
\[\max_{\pi_\theta} \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi_\theta (y \vert x)} [r_\phi (x, y)] - \beta \mathbb{D}_\textrm{KL} [\pi_\theta (y \vert x) \;\|\; \pi_\textrm{ref} (y \vert x)]\]이처럼 제약 조건을 추가해 모델이 리워드 모델의 분포에서 멀리 벗어나는 것을 방지하고 학습 불안정성도 조절합니다. 이 같은 형태로 PPO를 통해 학습하는 것이 기존의 RLHF입니다.
이제 DPO에서는 앞서 loss 수식을 다음과 같이 변형할 수 있습니다.
수식을 매우 잘 정리한 분3이 있어 가져왔습니다. 보다 상세한 내용은 해당 블로그를 참고하시기 바랍니다. 중간 과정은 생략하고 리워드 모델을 배제한 형태로 수식을 유도하면 다음과 같이 DPO 수식을 정리할 수 있습니다.
\[\mathcal{L}_\textrm{DPO} (\pi_\theta; \pi_\textrm{ref}) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \bigg[ \log \sigma \bigg( \beta \log \frac{\pi_\theta (y_w \vert x)}{\pi_\textrm{ref} (y_w \vert x)} - \beta \log \frac{\pi_\theta (y_l \vert x)}{\pi_\textrm{ref} (y_l \vert x)} \bigg) \bigg]\]RLHF에서 리워드 모델인 \(r_\phi\)를 두지 않고 수식 전개로 새로운 형태의 loss를 도출했습니다. 물론 RLHF 과정을 단순 수식 한 줄로 완전히 정의할 수 있느냐는 별도의 문제이지만, 이처럼 유사한 성능을 내면서도 훨씬 더 단순한 형태로 새로운 수식을 정의했다는 점은 무척 인상적입니다.
이제 DPO는 다음과 같은 형태로 그래디언트를 구하고 업데이트하면 됩니다.
논문에 제시된 대로 정답Chosen과 오답Rejected과의 가능성, 즉 마진을 최대화하면서 정답의 가능성은 높이고, 오답의 가능성은 낮추는 방향으로 업데이트합니다.
코드
그렇다면 실제 코드 구현을 살펴보겠습니다.
허깅페이스의 DPOTrainer 구현으로 완전한 형태를 참고할 수 있고, 먼저 PyTorch로 from scratch 버전을 구현한 분이 있어 이를 통해 좀 더 쉽고 직관적으로 살펴보도록 하겠습니다.
prompt_preferred_ids는 prompt + chosen을 max_length 만큼 늘여서 붙이고 나머지 공간은 모두 패딩으로 채웁니다.prompt_dispreferred_ids도 동일합니다. 단, rejected와 붙입니다.- 모델을 통과하고 model_preferred_logits를 구합니다. 기본 설정으로
[batch_size, 1024, vocab_size]가 됩니다.get_log_prob()함수에서 평균 log probability(이하 logp)를 구합니다. 먼저 선택 토큰의 logp를 구합니다. 사이즈는[1024]입니다. reponse_mask는 prompt 토큰에 0이고 나머지는 모두 1로 되어있는데, 패딩도 1로 처리되어 있어 이상해 보여 DPOTrainer와 직접 비교해봤습니다. DPOTrainer에서는 패딩이 제외처리 되어 있으며, 따라서 여기서도 패딩이 제외되는게 맞습니다. 관련하여 버그 리포팅을 진행했습니다.return response_log_probs / response_lengths이므로 길이만큼의 평균 logp를 리턴합니다. 마찬가지로 패딩 처리에 버그가 있습니다.
- 결과적으로
model_preferred_logprob는 프롬프트를 제외한 모든 토큰 로그 확률의 평균입니다. model_dispreferred_logprob는 프롬프트를 제외한 rejected에 대한 모든 토큰 로그 확률의 평균입니다.- 레퍼런스 모델도
ref_preferred_logprob,ref_dispreferred_logprob를 동일하게 구합니다. - 이제
calculate_DPO_loss()함수에서 실제로 dpo loss를 구현합니다.preferred_relative_logprob는model_preferred_logprob - ref_preferred_logprob입니다. 즉, 상대 logp입니다.dispreferred_relative_logprob도 동일합니다.reward_accuracies = preferred_relative_logprob > dispreferred_relative_logprob입니다. preferred의 logp가 더 크면 1이 됩니다.reward_margins = preferred_relative_logprob - dispreferred_relative_logprob입니다. margin은 말 그대로 logp의 차이입니다. 클수록 좋겠죠.- 최종적으로
loss = -F.logsigmoid(beta * (preferred_relative_logprob - dispreferred_relative_logprob))입니다. 수식에서는 log(pref/dispref)인데 여기서는 단순화하여 pref - dispref로 구현되어 있습니다. DPOTrainer에서도 동일하게logits = logratios - ref_logratios로 단순하게 처리하고 있습니다. 물론 실제 값은 수식과 다르지만 경향성은 동일할 것입니다. 다만 DPOTrainer에선 loss 계산 시 다음과 같이 수식에는 없는label_smoothing을 추가로 활용합니다.
logits = (chosen_logps - rejected_logps) - (ref_chosen_logps - ref_rejected_logps)
losses = (
-F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)
- F.logsigmoid(-self.beta * logits) * self.label_smoothing
)
실행
DPO 실행 시 다음과 같은 패턴을 보입니다.
loss가 줄고 있고, reward margin도 점점 커지고 reward accuracy는 거의 1에 수렴하므로 제대로 학습되고 있는 것으로 보입니다. rejected tokens인 dispreferred relative logprob도 잘 줄어들었으나 한 가지 문제는 preferred relative logprob도 함께 줄어들었다는 점입니다. 물론 reward margin이 더 커졌으니 별 문제는 없지만 rejected에도 실제로 chosen으로 나올법한 토큰이 많이 포함되어 있기 때문에 이에 따라 preferred logp도 함께 줄어드는 것으로 보입니다. 애초에 chosen과 rejected의 변별력이 매우 높다면 이 같은 현상이 감소할 것으로 보입니다.
허깅페이스 DPOTrainer의 metric은 다음과 같습니다.
rewards는 레퍼런스 모델과의 차이입니다. 왼쪽 맨위부터 rejected 차이 평균, chosen_rewards - rejected_rewards 즉 마진 평균, chosen 차이 평균, accuracies는 chosen_rewards > rejected_rewards의 평균입니다.
여기는 모델의 출력값입니다. 레퍼런스 모델과의 차이가 아님에 유의하세요. 각각 모델 logp의 평균, logits의 평균입니다.
