[yongggg's] Direct Preference Optimization: Your Language Model is Secretly a Reward Model; DPO Review

안녕하세요 이번 장에서는 RLHF의 대안으로 떠올랐던 DPO에 대해서 살펴볼까 합니다.

현재는 ORPO, TPO, KTO, SimPO 등등 많은 RLHF 대안의 알고리즘이 많이 나왔는데, 제일 먼저 나왔던 DPO를 살펴본 뒤 다른 것들도 DPO와 무엇이 다른지의 관점으로 보려고합니다.

DPO는 PPO, TRPO와 같은 강화학습 policy를 사용하지 않습니다. 따라서, reward model이 따로 필요하지 않는데, reward 함수와 policy만을 매핑하여 human preference data에 최적화할 수 있다고 합니다.

DPO 알고리즘을 알아보기 전에, RLHF부터 간단하게 살펴보면서 논문 리뷰 해보겠습니다!

 

1. RLHF

RLHF는 LLM이 생성하는 문장들을 사람의 선호에 맞게 다시한번 Tuning을 하면, 같은 input이어도 사람의 선호도에 가깝도록 문장을 생성하도록 유도한다. 이 과정은 SFT 이후에 RLHF로 마무리하는 작업으로 많이 사용되며, ChatGPT, LLAMA2 등의 모델도 RLHF로 모델을 완성했다.

RLHF는 인간의 선호도(human preference) dataset을 이용하여 Reward model을 훈련하고, 이를 생성 결과에 반영하도록 한다. objective는 1) 높은 reward를 주는 데이터에 가깝게 학습할 수 있도록 LLM을 최적화 하며, 2) LLM이 기존 return 값의 Policy와는 너무 멀어지지 않도록 제약을 주는 것으로 구성된다. (이때, 제약은 KL divergence를 이용함.) 

위의 그림의 왼쪽처럼 reward 모델과 기존 언어 모델의 policy를 잘 align할 수 있어야 하는 구조이다. 이 학습 단계는 크게 3 단계로 구성된다. 

1) Supervised Fine-Tuning (SFT): labeled dataset으로 학습 데이터셋을 Fine-tuning 하는 과정이다. (생성 Task)

2-1) Reward Modeling Phase : 다음으로 input이 SFT에 들어갔을 때, 나온 여러 개의 output들에게 점수를 매긴다. 인간이 가장 선호하는 output이 최고점을 받고, 그렇지 않으면 낮은 점수를 받도록 reward를 설정한다.

 

ex) 데이터 구축: $x$라는 prompt를 모델의 input이라고 가정하자, SFT로 학습한 모델로 2개의 output을 생성하여 output $y1$, $y2$를 얻었다고 하자. 이를 수식으로 표현하면 다음과 같다.

$$ (y_{1}, y_{2}) \sim \pi^{SFT}(y | x) $$ 

이제 인간이 등장하여, $y_{1}$과 $y_{2}$ 중 어떤 것이 더 선호되는 output인지 점수를 매긴다(reward를 준다). $y_{1}$과 $y_{2}$ 중 인간이 선호하는 답변을 $y_{w}$, 덜 선호하는 답변을 $y_{l}$라고 하자.

모델링: 원래는 여러 개의 sample이 있으며, 이러한 데이터로 reward model 학습하여, reward model이 선호하는 답변을 낼 수 있도록 학습 되어야한다. 이 latent reward model을 $r(y,x)$라고 표현하자 $r$을 정의하는 데에는 여러 알고리즘이 있지만, 주로 BT(Bradly-Terry) 방법을 사용하며, 수식은 아래와 같다. $y_{1}$의 reward가 더 높은 확률을 확률 모델로 정의한 것이다.

$$ P^{*}(y_{1} > y_{2} | x) = {exp(r^{*} (x, y_{1})) \over exp(r^{*} (x, y_{1})) + exp( r^{*} (x, y_{2}))} $$

 

이 수식을 통해 학습 Object를 정의할 수 있다. $p(y_{1} > y_{2} | x)$는 $y_{1}$의 reward가 $y_{2}$의 reward보다 높은 확률이다. 이는 인간이 $y_{1}$에게 $y_{2}$보다 높은 reward를 부여했다면, $p(y_{1} > y_{2} |x)$가 $p(y_{2} > y_{1} |x)$보다 높은 확률을 가지길 바랄 것이다. 이 생각을 토대로 NLLoss를 다음과 같이 정의할 수 있다. ($y_{w}$: 선호 답변, $y_{l}$: 비선호 답변)

$$ L_{R} (r_{\phi}, D) = -E_{(x,y_{w}, y_{l}) \sim D} [log \sigma (r_{\phi} (x, y_{w}) - r_{\phi} (x,y_{l}))] $$

여기서 reward model의 weights $r_{\phi}$의 초기 세팅은 SFT 모델의 마지막 linear layer output과 같다. human feedback을 받은 데이터로 $r_{\phi}$를 업데이트 한다.

 

2-2) RL Fine-Tuning Phase: 위의 reward loss로만 학습한다면, reward model이 생성한 문장과 현재 SFT 모델이 생성한 최적 답변이 상이할 수 있다. 다라서, reward loss에 constraint를 추가한다.

$$ \max\limits_{\pi_{\theta}} E_{x \sim D, y \sim \pi_{\theta} ( y|x)} [ r_{\phi} (x,y) ] - \beta D_{KL} [ \pi_{\theta} (y| x) || \pi_{ref} (y |x )] $$

KL divergence는 두 확률이 같을 수록 0에 수렴한다. 새로 학습한 $\pi$모델과 기존 initial model $\pi_{ref}$(SFT) 모델 간의 큰 차이가 없다면, constraint는 0에 가까워질 수 있다. 이 KL term을 이용하여 언어 모델이 너무 reward model 쪽으로 치우치지 않도록 할 수 있다. 최종 loss는 미분이 되지 않기 때문에, 강화학습 방법론인 PPO를 통해 학습한다. 최종 Loss를 강화학습에서 학습해야할 reward로 바라보고 최적화 하는 형태인 것이다.

 

2. DPO

DPO는 위에서 실행한 '최종 loss를 reward로 하여 PPO 알고리즘의 업데이트'를 거치는 과정을 없앤다. RLHF의 reward modeling phase를 바로 언어모델을 학습하는 phase로 변형한다. 그래서 reward modeling을 preference modeling 과정으로 치환환다. (reward는 강화학습에서 사용하는 단어이기에 DPO에서는 강화학습을 하지 않는다는 것을 강조하기 위해 reward를 preference로 치환한다.)

 

DPO 수식 유도

2-1. optimal policy $\pi$ : 위에서 설명한 RLHF의 RL Fine-Tuning Phase에서 구한 최종 Loss를 변형한다.

$$ \begin{align} \max\limits_{\pi} E_{x \sim D, y \sim \pi} &[r(x,y)] - \beta D_{KL} [\pi (y|x) || \pi_{ref} (y|x)] \tag{1} \\ &=\max\limits_{\pi} E_{x \sim D} E_{y \sim \pi (y|x)} \left[ r(x,y) - \beta log{\pi (y|x) \over \pi_{ref}(y|x)} \right] \tag{2} \\      &=\min\limits_{\pi} E_{x \sim D} E_{y \sim \pi (y|x)} \left[ log {\pi (y|x) \over \pi_{ref} (y|x)} - {1 \over \beta} r(x,y) \right] \tag{3} \\  &=\min\limits_{\pi} E_{x \sim D} E_{y \sim \pi (y|x)} \left[ log {\pi (y|x) \over {1 \over Z(x)} \pi_{ref} (y|x) exp \left( {1 \over \beta} r (x,y) \right) } - log Z(x) \right] \tag{4} \end{align}$$

여기서 $\pi (y|x)$는 업데이트할 모델, $\pi_{ref} (y|x)$는 기준 모델(SFT)이며, (1)에서 (2)로 갈때 $\beta$를 낀 항에서 $KL(q || p) = q log q - q log p $의 정의가 사용되었으며, (2)에서 (3)으로 갈때, $-{1 \over \beta}$를 곱하면서 $ \max $ 가 $ \min $ 문제로 바뀌었다. (4) 식에서 결국 분모의 ${1 \over Z(x)} \pi_{ref} (y|x) exp \left( { 1 \over \beta} r(x,y) \right) \tag{5}$ 항을 maximize 하는 문제가 된다.

또 $Z$는 아래의 식과 같다. 업데이트 중 모델 $\pi$와 무관한 partition function이다. $\pi_{ref}$와 x로 이루어져 있지만, $\pi$와는 무관하다.

$$ Z(x) = \sum\limits_{y} \pi_{ref} (y|x) exp \left( {1 \over \beta} r(x,y) \right) $$

 

위의 과정을 통해, 새로운 optimal policy $\pi^{*}$를 다음과 같이 쓸 수 있다. 이는 DPO 수식의 (5) 부분과 같다. 이 term을 maximize하면, 전체 Loss를 minimize하는 것과 같기 때문이다.

$$ \pi^{*} (y|x) = {1 \over Z(x)} \pi_{ref} (y|x) exp \left( {1 \over \beta} r(x,y) \right) $$

이 $\pi^{*}$를 DPO 수식 (4)에 다시 대입해보자. 그러면, $\min E$ 내부 term이 $ log {\pi (y|x) \over \pi^{*} (y|x) } - log Z(x)$가 된다. $E \pi (y|x)$이기 때문에, 수식을 풀면 $KL(\pi | \pi^{*})$이 된다. $Z$는 $\pi$와 관련이 없고, 우린 $\pi$를 최적화하기 위한 수식을 구하는 것이 목적이기 때문에 $Z$와 관련된 term은 날아간다. 그러면, $\pi$가 $\pi^{*}$에 근접할 때, loss가 가장 작다는 의미가 된다.

따라서 $\pi$와 $\pi^{*}$는 아래와 같다.

$$ \pi (y|x) = \pi^{*} (y|x) = {1 \over Z(x)} \pi_{ref} (y|x) exp \left( {1 \over \beta} r(x,y) \right) $$ 

 

2-2. reward model 변경

2-1. 에서 도출한 것은 다음 식 (6)과 같다.

$$ \pi_{r} (y|x) = {1 \over Z(x)} \pi_{ref} (y|x) exp \left( {1 \over \beta} r(x,y) \right) \tag{6} $$

식 (6)에서 양변에 $log$를 취한 후, reward $r(x,y)$를 좌변으로 옮기면 아래와 같이 식(7)을 도출할 수 있다.

$$ r(x,y) = \beta log {\pi_{r} (y|x) \over \pi_{ref} (y|x)} + \beta log Z(x) \tag{7} $$

이제 RL Modeling phase에서 사용했었던 Bradley-Terry model의 $r$과 $r^{*}$ 자리에 각각 식 (7)을 대입한다. 이제 DPO 만의 새로운 Preference model을 도출한 것이다.

$$ p^{*} (y_{1} > y_{2} | x) = {1 \over 1 + exp \left( \beta log {\pi^{*} (y_{2} |x) \over \pi_{ref} (y_{2} | x) } - \beta log {\pi^{*} (y_{1} |x ) \over \pi_{ref} (y_{1} |x) } \right)} \tag{8} $$

 

RLHF 단계에서와 같이 $y_{1}$, $y_{2}$를 human feedback을 통해 선호 및 비선호 태그를 붙인 $y_{w}$, $y_{l}$로 식 (8)에 대입한 후 NLL로 나타내면, 아래와 같은 최종 DPO loss가 도출된다. 이 식에서는 따로 reward model function(RLHF에서 $r_{\phi}$) 없이, 직접적으로 언어 모델의 loss를 도출했다.

$$ L_{DPO} (\pi_{\theta} ; \pi_{ref} ) = -E_{(x,y_{w}, y_{l}) \sim D} \left[ log \sigma \left( \beta log {\pi_{\theta} (y_{w} |x) \over \pi_{ref} (y_{w} |x)} - \beta log {\pi_{\theta} (y_{l} |x) \over \pi_{ref} (y_{l} | x) } \right) \right] $$

3. Theoretical Analysis of DPO

DPO에서 reward model이 필요없다는 주장을 뒷받침하는 근거는 '언어 모델은 이미 reward model의 역할을 하고 있다.'이다. 이를 증명하기 위해서는 'reward 함수가 언어 모델 policy로 충분히 도출될 수 있다.'를 밝혀야 한다. 이를 밝히기 위해, 두 가지의 명제가 필요하다.

명제 1. Bradley-Terry Model에 근거한 두 reward 함수는 같은 클래스 (c)에서 같은 preference distribution을 따른다.
$ \rightarrow P_{\pi} (c) = P_{\pi_{ref}}(c)$

명제 2. 강화학습에서 같은 클래스 (c)에 대한 두 reward function의 optimal policy는 같다.
$ \rightarrow r(x,y) = \pi_{ref} r(x,y) $

정의 1. $r(x,y)$와 $r^{'}(x,y)$가 어떤 함수 $f(x)$에 대해 $r(x,y) - r^{'}(x,y) = f(x)$를 만족하면, $r$과 $r^{'}$는 동일하다.

 

명제 1은 probability를 단어 token $y_{1}, ..., y_{k}$에 대한 랭킹 $( \tau )$를 reward라고 생각하자. probability를 softmax로 가정하고, 정의 1을 사용해 $r^{'}(x,y)$ 자리에 $r(x,y) + f(x)$를 대입하면 쉽게 증명이 가능하다. 이를 통해 probability를 계산하는 중에 어떤 식이 생략될 수 있다는 것을 알 수 있다. 즉, preference model에서 constraint를 주입해도 괜찮다라는 결론이 나올 수 있다.

$$ \begin{align} P_{r^{'}} (\tau | y_{1}, ..., y_{K}, x) &=  \prod\limits^{K}_{k=1} {exp(r^{'}(x, y_{\tau (k)})) \over \sum\limits^{K}_{j=k} exp (r^{'}(x,y_{\tau (j)}))} \\ &= \prod\limits^{K}_{k=1} {exp(r(x, y_{\tau (k)}) + f(x)) \over \sum\limits^{K}_{j=k} exp (r(x,y_{\tau (j)}) +f(x)) } \\ &= \prod\limits^{K}_{k=1} {exp(f(x)) exp(r(x,y_{\tau (k)}) \over exp(f(x)) \sum\limits^{K}_{j=k} exp(r(x,y_{\tau (j)}))} \\ &= \prod\limits^{K}_{k=1} {exp(r(x,y_{\tau (k)})) \over \sum\limits^{K}_{j=k} exp(r(x,y_{\tau (j)} )) } \\ &= p_{r}(\tau | y_{1}, ..., y_{K}, x) \end{align}

 

명제 2는 동일한 클래스에서 모든 reward function들이 하나의 optimal policy를 가진다는 의미이다. 이또한 $r^{'}(x,y) = r(x,y) + f(x) $를 이용하면 아래와 같이 쉽게 도출한다.

$$ \begin{align} \pi_{r^{'}}(y|x) &= {1 \over \sum_{y} \pi_{ref} (y|x) exp \left( {1 \over \beta} r^{'} (x,y) \right)} \pi_{ref} (y|x) exp \left( {1 \over \beta} r^{'} (x,y) \right) \\ &= {1 \over \sum_{y} \pi_{ref} (y|x) exp \left( {1 \over \beta} (r(x,y) + f(x)) \right)} \pi_{ref} (y|x) exp \left( {1 \over \beta} (r(x,y) + f(x)) \right) \\ &= {1 \over exp \left(  {1 \over \beta} f(x) \right) \sum_{y} \pi_{ref} (y|x) exp \left( {1 \over \beta} r(x,y) \right)} \pi_{ref} (y|x) exp \left( {1 \over \beta} r(x,y) \right) exp \left( {1 \over \beta} f(x) \right) \\ &= {1 \over \sum_{y} \pi_{ref} (y|x) exp \left( {1 \over \beta} r(x,y) \right)} \pi_{ref} (y|x) exp \left( {1 \over \beta} r(x,y) \right) \\ &= \pi_{r} (y|x) \end{align} $$

 

이제 위의 두 명제를 통해 reward model $r(x,y)$가 모델 policy인 $\pi_{r}$을 통해 표현할 수 있음을 보인다. (RLHF는 $\pi_{r}$을 최적화하기 위해 PPO를 사용하지만, DPO에서는 $\pi_{r}$을 단지 언어모델의 policy $\pi$로 표현할 수 있었다.) DPO 설명에서 식(6)을 가져와서 다시 살펴보면, 이때 $\pi_{r}$은 DPO의 최적의 policy이자 reward function을 의미했다.

$$ \pi_{r} (y|x) = {1 \over Z(x)} \pi_{ref} (y|x) exp \left( {1 \over \beta} r(x,y) \right) \tag{6} $$

$r(x,y)$와 $\pi_{r}$에 대해 $f(x) (x,y) = r(x,y) - \pi_{r} (y|x)$가 성립한다면(정의 1), $r$과 $\pi$에 대한 두 명제도 성립할 것이다. 임의의 projection $f$에 대해 아래의 식을 가정한다. $(f \fallingdotseq \beta log Z)$

$$ f(r; \pi_{ref}, \beta) = r(x,y) - \beta log \sum_{y} \pi_{ref} (y|x) exp \left({1 \over \beta} r(x,y) \right) \tag{7}$$

여기서 위 식을 만족하는 $f$가 존재하는 지만 안다면, $r(x,y)$와 모델 policy $\pi$가 동일하다고 할 수 있다. 그리고 만족하는 $f$는 반드시 있다. 식 (6)을 다시 한 번 보면, DPO reward model 증명 식(7)과 닮았다. 즉, $f$는 $\pi$에 대한 식이다. 어떠한 $f$에 대해 DPO reward model 증명 식(7)이 성립함을 보였다.

$$ f(r; \pi_{ref}, \beta)(x,y) = \beta log {\pi_{r} (y|x) \over \pi_{ref} (y|x) } $$

이로서 언어모델은 그 자체로 reward model의 역할을 하고 있음을 밝혔다. 추가적으로, $f$ 자체로도 $r(x,y)$의 역할을 할 수 있다.

reward 함수의 $log$값에 $Z$라는 constraint를 더한 것 (명제 1에 의하면 constraint를 임의 주입해도 reward 함수의 probability는 동일했다.)이기 때문이다.

4. Experiment

시실험 내용 중에 유심히 볼 것 들이 있다. 첫쨰로 PPO(RLHF)를 적용한 모델은 KL이 일정이상 커지면 해당 policy에 reward가 아예 계산되지 않지만, DPO를 적용한 모델은 KL에 independence하게 reward를 주고 있음을 볼 수 있다.

그리고 RLHF에 비해 DPO는 temperature에 대해 robust 하게 보여진다.