[RLHF] KTO: Model Alignment as Prospect Theoretic Optimization

 KTO: Model Alignment as Prospect Theoretic Optimization

 

link : https://arxiv.org/abs/2402.01306

 

 

 

 

Kahneman & Tversky’s prospect theory 에 따르면 사람의 편향되어 있지만 잘 정의된 방식으로 random variable을 인식한다고 한다. (예를 들어서 사람은 손실을 싫어한다.) 우리가 LLM을 human feedback에 aligning하려는 목표는 이러한 편향이 암묵적으로 포함되어 있다고 볼 수 있다. DPO는 cross entropy 최소화 objective를 통해 부분적으로 이를 가능하게 했지만 현재 DPO와 같은 utility function은 Kahneman & Tversky’s prospect theory가 이야기하는 것과는 다르다. 따라서 해당 논문에서는 Kahneman & Tversky Model을 이용해 preference의 log-likelhood를 최대화하는 것 대신 utility of generation을 직접적으로 최대화 하는 방법을 제안한다. 기존 RLHF에서는 preference dataset (어떤 답변을 더 선호하는지 포함된 데이터셋)이 필수적이었으나 KTO의 가장 큰 맹점은 이 preference 데이터셋이 필요하지 않다는 것이다. 단순히 해당 답변이 desired인지 아닌지만 필요하다. 즉, 선호 pair가 필요 없다는 것이다.

 

 

✲ Introduction

 

ChatGPT, DPO 등의 등장으로 우리가 LLM을 단순히 SFT하는 것 보다 human feedback에 aligning하는 RLHF 기법이 훨씬 효용적이라는 것이 밝혀졌다. 하지만, human feedback은 real word에서 수집하기 어렵고 희소함에도 일종의 선호도 맥락에서만 해당 기법이 논의된다. 그리고 이러한 alignment method가 왜 잘 작동하는지, preference 형태의 feedback이 필요한지 아닌지, Kahneman & Tversky’s prospect theory의 관점에서 설명한다. 해당 이론은 인간은 expected value를 최대화하지 못하는 불확실한 event에서도 의사결정을 내리는지 설명한다. 예를 들어 어떤 기전점에 의해 사람은 이익보다는 손실에 더 민감하게 반응하는 속성을 통해 사람은 편향적이지만, 잘 정의된 무작위 변수를 인지한다 라는 것을 공식화한다.

 

저자들은 더미 +1/-1 보상에 대해 offline-PPO을 실행하면 대부분의 규모에서 DPO의 성능과 일치 시킬 수 있다는 것을 발견하였고 이는 loss에 inductive bias가 충분하다면 preference dataset이 필요하지 않을 수 있음을 시사한다.

 

 

‣ KTO의 contribution

 

1. KTP는 1B부터 30B까지의 모델에서 DPO와 일치하거나 능가하는 성능을 보였다. 이는 DPO의 n pair를 KTO에서는 2n개의 데이터셋으로 늘릴 수 있고 pair set보다 더 약한 signal을 학습함에도 더 좋은 generation을 보여주며 이에 대한 theoretical explanation도 다룬다.
2. KTO는 극단적으로 imbalance한 데이터셋으로도 DPO와 비슷한 성능을 보여줄 수 있다. 따라서 선호셋이 human alignment에 기인하지 않는 다는 것을 의미한다.
3. 또한 pretrained-model의 성능이 충분히 좋다면 KTO는 성능 손실 없이 SFT과정을 스킵할 수 있다. 반면 DPO는 모든 스케일에서 SFT가 없는 경우 성능이 악화되었다.

 

✲ Background

‣ RLHF

 

데이터셋에 인풋 x 그리고 두 아웃풋 y_w, y_l이 있을 때 각각 선호하는 답변, 선호하지 않는 답변이 있다. 여기서 선호에대한 r*는 true reward function이다. 이는 y_w이 y_l보다 선호된다는 것을 specific function class, 일반적으로 bradley-terry model, 가 포착한 확률이다. 

 

sigmoid는 logistic function

 

사람의 true reward는 얻기 어려우므로 rewrd model r_phi는 proxy learning을 하게 되는데, human preference data의 negative log-likelihood를 minimizing하는 방향으로 학습한다.

 

하지만 보상을 최대화 하는 것 만으로는 문법적인 텍스트생성과 같은 것들을 희생한다. 이를 피하기 위해 reference model pi_ref와 너무 멀지 않게 하는 KL divergence 텀을 도입한다. 

 

 

해당 objective는 not differentiable하기 때문에 해당 objective를 optimize하기 위해서는 PPO와 같은 RL 알고리즘이 사용된다.

그러나 RLHF는 주로 generations을 샘플해야 하기 때문에 학습이 느리고 특히 분산환경에서 unstable 하다.

이러한 이유로 최근 연구에서는 preferred, dispreffered generation 의 margin을 키우는 closed-form losses를 고안하는데 초점이 맞춰져 있다. (ex. SLiC, DPO ..)

 

 

그리고 DPO는 특히 RLHF와 mathematical equivalence로 인해 popular하다.

 

 

✲ A Prospect Tehoretic View of Alignment

 

카네만 트레베르스키의 전망 이론은 불확실한 사건에 직면한 인간은 왜 expected value를 극대화하는 쪽으로 의사결정을 내리지 않는지 설명한다. 예를 들어 인간은 손실 회피 성향이 있기 때문에 80%확률로 100달러를 얻고 20%확률로 0달러를 얻는 gamble이 주어지면 60달러의 확실한 값을 선택할 수 있다.

 

 

‣ Prospect Theory

 

전망 이론은 human utility는 Value function과 Weighting function에 따라 달라진다.

 

Definition 3.1:

Value function v: z -> R은 어떤 z_ref에 상대적인 결과 z를 perceived(or subjective) value에 매핑한다. 예를 들어 이 함수는 사람이 같은 크기의 상대적 이득보다 상대적 손실에 더 민감하게 반응하는 경향이 있다는 사실을 포착한다.

 

Definition 3.2:

가중치 함수 w는 누적 확률을 perceived 누적 확률에 매핑하는 capacity function의 도함수이다. 이 함수는 예를 들어 인간이 희귀한 사건의 발생 가능성을 과대평가하는 경향이 있다는 것을 포착한다. w_z는 z에 부여된 가중치를 의미한다.

 

Definition 3.3:

Reference point는 input-output sampled pair에 관한 보상 기댓값

 

random variable Z의 utility는 

 

 

해당 결과의 함수이다. 하지만 사람은 LLM 전체 확률 분포를 알 수 없으므로 weight function은 중요하지 않고 value function에 초점을 맞춘다. 실제 인간에게 gamble을 제시하고 그에 상흥하는 확실성 묻는 실험을 통해 카네만 트레베르스키는 인간의 가치에 대해 아래와 같은 함수 형태를 제안했다.

 

 

 

여기서 알파 0.88, 감마 2.25는 중앙값이고 알파는 utility의 변화 속도를, 감마는 손실 회피 정도를 제어한다. 카네만 트레베르스키 value function은 중앙 값의 모양은 figure 2에 표현되어 있고 개인마다 다르다는 것에 유의해야 한다.

 

 

✲ Kahneman-Tversky Optimization

 

더미 +1/-1 리워드를 사용한 Offline PPO의 놀라운 성공은 올바르게 HALO(Human-aware loss function)을 사용하면 단순 이진 signal만으로 DPO 수준의 성능에 도달할 수 있음을 시사한다.

보다 원칙적인 접근 방식으로 저자들은 Kahneman-Tversky model of human utility를 사용한 HALO를 유도하고, 이는 선호 log-likelihood를 최대화하는 것 대신 직접적으로 utility(human feedback기반 보상)를 최적화 한다. 이 Kahneman-Tversky Optimization (KTO) loss는 오직 desired/undesired 두 시그널만 필요하므로 현실세계에서 더 수집하기 쉽다.

 

‣ Derivation

 

DPO에서 RLHF와 equivalence를 갖는 유도를 통해 optimal policy pi*와 true reward r*는 위와 같이 유도되었다. (https://ebbnflow.tistory.com/382)

여기에 DPO에서는 Bradley-Terry model을 연결하고 negative logarithm을 취해 DPO objective를 도출하였다. KTO는 이것 대신 human utility에 대한 Kahneman-Tversky model를 연결하고 몇가지 추가 변형을 적용하여 objective를 도출한다.

 

1. Kahneman-Tversky Value function의 지수 알파는 최적화를 어렵게 만드므로 v_KTO를 gain에서 concave이고 loss에서 donvex인  logistic function으로 설정한다. 그리고 손실 회피 파라미터 gamma_D, gamma_U를 사용해 desired/undesired loss에 대해 각 가중치를 부여한다.

 

2. Kahneman-Tversky Value function은 사람의 금전적인 gamble을 대상으로 한 실험으로 도출된 것이고, LLM에서의 보상은 금전과 관련이 없으므로 RLHF의 objective에 따라 implicit reward r_KTO로 설정한다.

3. 하나의 선호되지 않는 generation만을 Reference point로 하는 대신, 사람은 자신이 본 모든 입출력 pair와 관련해 품질을 판단한다고 가정한다. 따라서 Reference point을 optimal policy를 따르는 x에 대해 설정하는 것이 아니고 any input에 대한 보상의 기댓값에 대해 설정한다. any input x'에 대한 partition function의 기댓값이 0이라는 가정으로 인해 pi*와 pi_ref의 KL divergence 텀에 beta를 곱한 것으로 간소화할 수 있다.

 

위와 같은 변경 사항을 적용하면 KTO objective는 아래와 같이 설계될 수 있으며 desired/undesired는 Kahneman-Tversky에서 gain/loss로 치환된다.

 

 

 

 

직관적으로, KTO는 만약 모델이 desirable example의 보상을 증가한다면 KL penalty 또한 증가함으로써 작동하여 loss에 대한 progress가 없기 때문에 작동한다. 이러한 강제는 모델이 무엇이 output을 desired하게 만드는지에 대해 학습하여 KL term을 유지하는(혹은 심지어 감소하게 만들어) 동안 보상이 증가할 수 있도록 한다. 그리고 또 KL 텀은 0보다 큰 값이기 때문에 -KL항은 항상 음수기 때문에 모델이 최적 솔루션에 가까워질 수록 보상이 빠르게 증가하는 것을 막아 stable 학습이 가능하도록 한다. 

 

 

✲ Implementation

 

 

실제로 KTO는 KL term을 estimate하여 사용하고, back-propagate에 사용하지 않고 보상을 조절하는데만 사용되서 학습을 stable하게 만든다. 즉 KL 텀은 순전히 loss가 staturate (어느 정도 범위에서 수렴) 정도를 조정하는 역할만 한다고 볼 수 있다.

 

그리고 beta는 DPO와 같은 역할을 하는데 클수록 reference model와 멀어지때 가해지는 패널티가 크다. 저자들은 이 값을 대부분의 데이터셋에서 0.1정도로 두는 것이 성능에 좋아서 0.1로 고정하였다고 한다.

 

 

그리고 n_D와 n_U는 각각 desirable, undesirable example의 갯수이고 gamma는 n_D, n_U의 비율에 따라 설정한 파라미터이다. 만약 1:1이었다면 gamma_U=1, gamma_D \in [1, 1.33] 정도로 설정하고, 9:1 정도였다면 gamma_U=1, gamma_D \in [10, 13.33]로 설정하였다.

 

 

---

 

‣ Review

 

offline-PPO가 더미 리워드에 대해서도 좋은 성능을 내는 것에 착안해, DPO에서 Bradley-terry 모델을 연결해 reward function을 없앤 RLHF equivalence으로 objective를 도출하는 것 대신, Kahneman-Tversky 이론을 적용해 휴리스틱한 objective를 디자인한 논문이다.

gamma라는 파라미터가 하나 더늘어나서 dataset에 따라 조정해줘야 하는 단점은 있지만 imbalance하고 약한 라벨이 있는 데이터셋에서도 사용할 수 있다는 장점이 있다.