Machine & Deep Learning

[yongggg's] KTO: Model Alignment as Prospect Theoretic Optimization review

Yonggg 2024. 10. 2. 13:41

안녕하세요 이번 장에서는 RLHF의 목적을 조금더 쉽게 이루기 위한 또 다른 알고리즘인 KTO에 대해 설명드리겠습니다.

KTO 연구는 Kahneman & Tversky’s prospect theory을 바탕으로 하며, KTO는 DPO와 유사하게 LLMs을 human feedback에 맞춰 조정하는 최정화 방법입니다. 이전 연구들은 현재의 선호도 기반의 Loss를 줄임으로써 기존의 Cross entropy 최소화 방법보다 효과적이었지만, KTO에서는 그런 선호도 기반의 방식 대신 인간의 효용을 직접적으로 최대화하는 방법을 사용하여 성능을 향상시킵니다.

KTO는 입력에 대한 출력이 바람직한지의 여부를 이진 신호로만 학습하여, DPO 등 기존 방법과 동일한 데이터를 사용하면서, 더 나은 성능을 보여준다고 합니다.

이제 부터 그 내용에 대해 자세히 알아보겠습니다!

Introduction

ChatGPT, DPO 등의 등장으로 우리가 LLM을 simple하게 SFT 하는 것 보다는 human feedback을 통해 alignment하는 RLHF와 같은 기법이 훨씬 효용적이라고 밝혀졌다. 어떻게 보면 human feedback은 현업에서 수집하기가 너무 어렵고 희귀한 데이터이지만, 그에 대한 결과는 어는 방법보다 가장 효과적이기 때문에 필요한 기법이다.

이 연구에서는 이러한 방법들이 왜 그렇게 잘 작동하는지, 그리고 feedback이 반드시 선호도의 형태여야 하는지를 이해하기 위해, Kahneman & Tversky’s prospect theory 관점에서 alignment를 재구성한다. 이 이론은 인간이 random variable 어떻게 편향되도록, 그러나 잘 정의된 방식으로 인식하는지 공식화한다. (ex: reference point에 비해 인간은 이득보다 손실에 더 민감하며, 이를 손실 회피라고 하는데, DPO, PPO-Clip과 같은 유명한 alignment 방법들이 이런 편향 중 일부를 내재적으로 모델링하고 있으며, 이는 데이터와 무관하게 그것들의 성공을 부분적으로 설명하는데 도움이 된다.)

그런 다음, 본 연구에서는 이러한 손실 함수를 보다 일반화한 형태인 Human-Aware Losses (HALOs)라는 새로운 손실 함수 class에 대해 제안한다.

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

background

- Pre-training

- SFT

- RLHF

https://yongggg.tistory.com/41

 

[yongggg's] Training language models to follow instructions with human feedback (InstructGPT) review

안녕하세요 이번에 설명 드릴 논문은 Open AI의 "Training language models to follow instructions with human feedback" 이라는 논문입니다. 최근 ChatGPT가 등장하면서, 이에 근간이 되는 이 InstructGPT 논문에도 관심이

yongggg.tistory.com

A Prospect Theoretic View of Alignment

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

Prospect Theory

전망 이론에서 인간의 효용은 가치 함수와 가중 함수에 의존한다:

Definition 3.1. A value function $v : Z \rightarrow \mathbb{R}$은 결과 $z$를 기준점 $z_{0}$에 상대적으로 매핑하여 지각된(perceived) or 주관적인 가치를 제공한다. 예를 들어, 이러한 함수는 동일한 크기의 이득보다 손실에 더 민감한 인간의 경향을 포착한다.

Definition 3.2. A weighting function $w$는 누적 확률을 지각된 누적 확률로 매핑하는 capacity function의 도함수이다. 예를 들어 이 함수는 인간이 드문 사건의 발생 확률을 과대평가하는 경향을 설명한다. 다시 말해, 인간이 희귀한 사건의 발생 가능성을 과대평가 하는 경향이 있다는 것을 포착한다. $w_{z}$는 $z$에 부여된 가중치를 의미한다.

Definition 3.3. Random Variable $Z$의 utility는 $u(Z) \triangleq \sum_{z \in Z} \omega_z v(z - z_0).$ 결과의 함수이다. 

하지만 사람은 LLM의 전체 확률 분포를 알 수 없기 때문에, weight function은 중요하지 않고, value function에 초점을 맞춘다. 실제 인간에게 gamble을 제시하고 그에 상응하는 확실성을 몯는 실험을 통해 카네만 트레베르스키는 인간의 가치에 대한 함수를 아래와 같이 제안했다.

$$ v(z, z_{\text{ref}}; \lambda, \alpha) = \begin{cases} (z - z_{\text{ref}})^\alpha & \text{if } z > z_{\text{ref}} \\ -\lambda(z_{\text{ref}} - z)^\alpha & \text{if } z < z_{\text{ref}} \end{cases} $$

해당 연구에서 $\alpha =0.88, \lambda =2.25$는 중앙값이고, $\alpha$는 utility의 변화 속도를, $\lambda$는 손실 회피 정도를 제어한다. 카네만 트레베르스키 value function은 중앙 값의 모양은 위의 figure 2에 표현되어 있고 개인마다 다르게 쓸 수 있다.

HALOs

Definition 3.4(HALOs). 모델의 학습 가능 파라미터를 $\theta$로 나타내고, $\pi_{\theta}: X \rightarrow P(Y)$는 aligned model, $\pi_{ref}$는 reference model. $l:Y \rightarrow \mathbb{R}^{+}$는 정규화 인자, 그리고 $r_{\theta}(x,y) = l(y)log \left[ {\pi_{theta}(y|x) \over \pi_{ref}(y|x)} \right]$는 implied(암시) reward라고 하자. 여기서 $Q(Y^{'}|x)$는 $Y$에 대한 reference point 분포를 나타내고, $v: \mathbb{R} \rightarrow \mathbb{R}$는 $(0, \infty)$ 범위에서 어디에서나 증가하고 오목한 함수이다. $(x,y)$에 대한 인간의 가치는 다음과 같이 표현된다.

$$v(r_{\theta}(x,y) - \mathbb{E}_{Q} [r_{\theta}(x,y^{'})])$$
함수 $f$는 $v$에 대한 human-aware loss(인간 인식 손실)이며, $a \in \{-1, +1\}$가 존재할 때 다음을 만족한다.

$$f(\pi_{\theta}, \pi_{ref}) = \mathbb{E}_{x,y \sim D} [a_{x,y} v(r_{\theta}(x,y) - \mathbb{E}_{Q}[r_{\theta}(x,y^{'})])] + C_{D}$$

여기서 $D$는 feedback data이며, $C_{D} \in \mathbb{R}$은 데이터에 특화된 상수이다.

고전적인 전망 이론 실험에서, $r_{\theta}$는 각 결과에 할당된 달러의 금액을 말한다. 여기서 $r_{\theta}$는 $\pi_{ref}$에서 $\pi_{\theta}$로 이동할 때, 조건부 surprisal이 감소한 정도를 나타내며, 이는 $l$에 따라 정규화되어 nats 단위로 측정된다. 이는 LLM의 사전 학습과 미세 조정에 사용되는 다음 Token 예측 목표에서 자연스럽게 도출될 수 있다. $\pi_{\theta}$가 정렬될수록, 바람직한 출력에 대해서는 $r_{\theta}$가 점점 더 긍정적으로, 바람직하지 않은 출력에 대해서는 점점 부정적으로 증가할 것으로 예상된다.

 

보상에 대한 또 다른 관점은 RLHF objecctive에서 도출된다. 이 목표를 최적화하는 정책은 closed form 표현을 가지며, 여기서 $Z(x)$는 분할 함수이다.

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

여기서 $l(\cdot)=\beta$로 설정하면,
$$ r^{*}_{\theta} = r(x,y) - \beta logZ(x)$$

$\theta^{*}$에서 HALO가 정의한 보상은 입력별 항에 의해 이동된 최적 보상일 뿐이며, 이는 $r^{*}_{\theta}$가 $r^{*}$와 동일한 동등 클래스에 속하며, 최적의 정책 $\pi^{*}$를 유도할 수 있음을 의미한다.

 

HALO에서 reference point는 인간의 관점에서 기대되는 보상이며, $Q(Y^{'} | x)$는 baseline을 구성하는 데 사용되는 예시들을 설명한다. 본 연구에서는 가치 함수가 gain(이득) 부분에서는 오목해야 하나, 반드시 손실 부분에서 볼록할 필요가 없다는 것을 요구한다. 이는 소수의 개인들이 이득과 손실 모두에 대해 위험 회피 성향을 보이기 때문에, 이 경우 가치 함수가 전체적으로 오목해지기 때문이다.

 

Theorem 3.5. DPO and PPO-Clip는 human-aware loss 이다.

손실함수 $f$가 HALO가 되기위해서는 인간의 가치 $v(r_{\theta}(x,y) - \mathbb{E}_{Q} [ r_{\theta} (x,y^{'})])$를 구성해야 한다. 여기서 $r_{\theta} (x,y) = l(x,y)log\pi_{theta}(y|x)$는 (정규화 인자 $l(y)$로 정규화된 implied 보상이며, $Q(Y^{'}|x)$는 입력 조건부 reference point 분포인 $\pi_{ref}(y|x)$, 그리고 $v: \mathbb{R} \rightarrow \mathbb{R}$는 전망 이론적 관점에서 $(0, \infty )$ 범위에서 오목하고 어떤 구간에서든지 비감소하는 value function이다.

DPO 손실은 다음과 같으며,

$$L_{\text{DPO}} (\pi_{\theta}, \pi_{ref}) = \mathbb{E}_{x,y_{w},y_{l}} \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]$$

여기서 $\beta >0$는 하이퍼파라미터이다. DPO는 다음과 같은 구성으로 HALO 조건을 만족한다. : $l(y) = \beta; r_{\theta}=\beta \log \pi_{\theta} (y|x) ; v(\cdot) = \log \sigma (\cdot) $은 어디서나 증가하고 오목한 함수이다. $Q$는 모든 질량(mass)를 $(x,y_{l})$에 두며, 여기서 $y_{l}$은 $x$에 대한 선호되지 않는 output으로 $ y > y_{l}$을 만족한다. 그리고 $a_{x,y} = -1$이다.

 

PPO-Clip loss는 해당 논문의 Appendix B참고.

Does being a HALO matter?

HALO와 non-HALO가 모두 다양한 함수 class이기 때문에 HALO 유용한지에 대한 답변은 어렵다. 그러나 본 연구는 HALO로 간주되는 popular 방법들과 그렇지 않은 방법들을 비교하여 이 질문에 답하려 한다.

  1. CSFT: Conditional SFT는 간단한 alignment 방법으로, 훈련 중 output에 control token을 앞에 추가하고, 추론 시에는 바람직한 생성을 유도하기 위해 입력에 해당하는 control token을 추가한다.(Korbak et al., 2023). 이는 HALO가 아니다. 왜냐하면, HALO는 $ - \log \pi_{ref}(y|x)$가 항상 reference point와 같아야 하기 때문이다. 그러나 reference point가 주어진 $x$에 대해 고정되기 때문에, $\pi_{ref}$가 non-uniform distribution일 때, 이는 불가능하다.
  2. SLiC: Sequence Likelihood Calibration (Zhao et al., 2023)은 선호도에 대한 최대 margin loss와 언어 모델링 손실을 결합한 방법으로, 위의 CSFT와 같은 이유로 HALO가 아니다.
    $$\begin{equation*} \begin{split} &L_{\text{cal}}(\pi_\theta) = \mathbb{E}_{x, y_w, y_l \sim D} \left[ \max\left( 0, \delta - \log \frac{\pi_\theta(y_w|x)}{\pi_\theta(y_l|x)} \right) \right] \\ &L_{\text{reg}}(\pi_\theta, \pi_{\text{ref}}) = \mathbb{E}_{x \sim D, y \sim \pi_{\text{ref}}(x)} \left[ - \log \pi_\theta(y|x) \right] \\ &L_{\text{SLiC}}(\pi_\theta, \pi_{\text{ref}}) = L_{\text{cal}}(\pi_\theta) + \lambda_{\text{reg}} L_{\text{reg}}(\pi_\theta, \pi_{\text{ref}}) \end{split} \end{equation*}$$
  3. DPO: DPO는 HALO임을 위에서 증명함. (Theorem 3.5.)
  4. PPO (offline): standard RLHF의 Objective는 주로 이전 단계 버전인 $\pi_{old}$로 부터 $\pi_{\theta}$가 얼마나 멀리 벗어날 수 있는지를 "clipping"하는 PPO-Clip으로 최적화된다:
    $$ L_{\text{PPO (offline)}} = - \mathbb{E}_{x,y,t \sim D} [ min(q_{\theta}A(x:y_{<t}, y_{t}), \\ clip(q_{\theta}, 1 - \epsilon, 1+ \epsilon) A(x:y_{<t}, y_{t}))] $$
    여기서 $q_{\theta} = {\pi_{\theta}(y_{t} | x: y_{<t}) \over \pi_{old}(y_{t} | x:y_{<t})}$이고, $A(x:y_{<t}, y_{t})$는 token 당 advantage 즉, 주어진 state에서 특정 token을 생성함으로써 얻는 잉여 이점을 말한다.

    PPO는 on-line 알고리즘이다. 생성된 sample은 현재 모델에서 추출되며, reward model에 의해 평가된 후 현재 버전을 업데이트 하는데 사용된다. 그러나 off-line 방법과의 공정한 비교를 위해 off-line 데이터만을 사용하는 PPO의 변형을 고려한다. 또한 RLHF는 분산 환경에서 매우 불안정하기 때문에, $\pi_{old}$를 업데이트하지 않고 $\pi_{ref}$로 유지하며, 전통적인 방식보다 덜 보수적으로 clipping을 적용한다. Baheti etal. (2023)은 이러한 변화와 전체 출력 시퀀스를 single action으로 처리하는 방법이 안정성을 크게 향상시킨다고 발견했다. 그러나 RLHF는 전통적으로 token-level의 advantages를 사용하기 때문에, 우리는 세 번째 변경 사항을 생략하고 첫 번째와 두 번째 변경 사항만 유지한다. PPO-Clip은 변하지 않고 HALO로 유지된다. (Appendix B.)
    지금 방법은 off-line이고 한 번만 업데이트를 진행하기 때문에, PPO라고 하는 것은 약간 부정확하다.  그러나 너무 많은 새로운 용어를 도입하지 않기 위해 이를 PPO(offline)이라고 부르겠다. 학습된 보상을 사용하는 대신, $y_{w}$와 $y_{l}$에 대해 +1 / -1의 더미 보상을 사용하여 더욱 간단하게 처리한다. 이 방법의 구현에 대한 추가 세부 사항은 Appendix C를 확인하길 바란다.

Experimental Setup

실험으로 Pythia-{1.4B, 2.8B, 6.9B, 12B}, Llama-{7B, 13B, 30B)를 기반으로, 다양한 규모의 LLM의 align 성능을 비교 분석하며, Mistral 및 그 것의 파생 모델에서도 수행된다. 모든 모델은 동일한 데이터와 설정 하에 학습했으며, 하이퍼파라미터만 모델에 따라 조정되었다.

주요 실험 설정

  • 학습 데이터는 Anthropic-HH, OpenAssistant, SHP등의 데이터를 사용했으며, GPT-4-0613을 이용해 모델의 output이 SFT target 보다 나은지 평가한다. 평가 기준은 helpfulness(도움이 되는지), harmlessness(해로운지), and conciseness(간결한지) 이다.
  • SFT target은 바람직한 출력이지만 반드시 최상의 output은 아니며, aligned model이 이를 능가할 수 있는 가능성이 있다. 그러나 대부분의 target은 사람이 작성한 것이기 때문에 매우 강력한 baseline으로 간주된다. 


주요 결과:

  • HALO 모델은 모든 scale에서 non-HALO 모델보다 성능이 좋거나 동등하며, 특히 13B 이상의 모델 크기에서 유의미한 차이 $(p < 0.05)$가 나타났다. HALO aligned Llama-13B, 30B 모델은 SFT target 품질을 50% 이상의 비율로 능가했다.
  • 7B parameter 이하에서, alignment는 SFT alone 대비 거의 성능 향상을 제공하지 못했다. 그러나 만약 기본 모델의 성능이 더 높거나 SFT 데이터 분포가 선호 데이터와 덜 유사했다면, alignment 단계에서의 성능 향상이 더 컸을 것이라고 지적한다.
  • 단순한 +1/-1 보상을 사용했음에도 불구하고, off-line PPO 변형이 Llama-30B 모델을 제외한 모든 모델에서 DPO와 동등한 성능을 보였다. 이는 Reward training에 큰 비중을 두는 기존의 통념에 도전하는 결과이며, 적절한 귀납적 편향을 가진 손실 함수에서 사용될 경우, 가장 단순한 보상도 유용할 수 있음을 시사한다. 이러한 성공에도 불구하고 본 연구의 off-line PPO baseline은 여전히 hyperparameter 민감성과 학습 불안정성을 겪고 있지만, traditional RLHF 만큼은 아니다.

Kahneman-Tversky Optimization

dummy +1/-1 보상을 사용한 PPO의 성공은 적절한 귀납적 편향이 주어지면, 좋음/나쁨을 나타내는 이진 신호만으로도 DPO 수준의 성능에 도달할 수 있음을 나타낸다. 비록 offline PPO 접근 방식 자체는 일정 규모를 넘어서지 못했지만, 이를 고려하면서 보다 원칙적인 접근법으로 우리는 이제 Kahneman-Tversky의 인간 효용 모델을 사용하여 HALO를 도출한다. 이를 통해 선호도의 Log Likelihood를 최대화하는 대신 직접적으로 효용을 최적화할 수 있게 된다.

Derivation

DPO에서 RLHF와 등가를 만족하는 유도를 통해 optimal policy $\pi^{*}$와 true reward $r^{*}$는 다음과 같이 유도되었다.

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

이 식을 가지고 DPO에서는 Bradely-Terry model에 대입하고, negative log likelihood를 취해 DPO objective를 도출했다. 

KTO는 이 대신, human utility에 대한 Kahneman-Tversky model을 연결하고 몇가지 추가 변형을 적용하여 objective를 도출한다.

 

Kahneman-Tversky Value function의 지수 a로 인해 최적화 과정에서 수치적 불안정을 겪기 때문에, 우리는 이를 gains에서는 오목하고, loss에서는 볼록한 성질을 가진 로지스틱 함수 $\sigma$로 대체한다.

위험 회피 정도를 제어하기 위해, 본 연구에서는 $\beta \in R^{+}$라는 hyperparameter를 가치 함수의 일부로 도입한다. $\beta$ 값이 클수록 값이 더 빠르게 포화 상태가 되며, 이는 인간이 이익 상황에서는 더 위험을 회피하고, 손실 상황에서는 더 위험을 감수하는 경향이 있다는 것을 의미한다. 실제로 이는 DPO 손실에서 $\beta$와 유사한 효과를 가지며, DPO에서는 $\pi_{\theta}$가 $\pi_{ref}$에서 얼마나 벗어나는지를 제어한다. DPO에서는 RLHF의 objective의 KL 제약이며, 보상의 일부로 작용한다. 다만 이 연구에서는 위험 회피를 명시적으로 제어하기 위해 도입되었다.

 

본 연구에서는 원래의 Kahneman-Tversky 가치 함수에서 손실 회피 계수 $\lambda$를 $\{ \lambda_{D}, \lambda_{U} \}$로 대체한다. 여기서 $\lambda_{D}$와 $\lambda_{U}$는 각각 바람직한 출력과 바람직하지 않은 출력에 대한 hyperparameter이다. 또한, 더 복잡한 방식으로 중요도에 대한 샘플링을 사용할 수 있다.

 

DPO에서처럼 하나의 비선호 출력을 reference point $z_{0}$로 사용하는 대신, 인간이 $y|x$의 품질을 간으한 모든 출력과 비교한다고 가정한다. 이는 $Q(Y^{'}|x)$가 정책임을 의미하며, reference point는 KL divergence $KL(\pi_{\theta}(y^{'}|x) || \pi_{ref}(y^{'}|x))$이다. 그러나 실제로는 다음과 같이 이 항에 대해 편향된 추정을 사용한다.

 

여기서 $\lambda_{y}$는 $\lambda_{D}(\lambda_{U})$을 나타내며 이는 각각 $y$가 desirable(undesirable)할 때를 나타낸다. 이 때, default KTO loss는 다음과 같다.

$$ L_{\text{KTO}} (\pi_{\theta}, \pi_{ref}) = \mathbb{E}_{x,y \sim D} [ \lambda_{y} - v(x,y)] \pi_{\theta} (y|x) $$

where,

$$ r_{\theta}(x,y) = \log \pi_{ref}(y|x) \\ z_{0} = KL(\pi_{\theta} (y^{'}|x) || \pi_{ref}(y^{'}| x))$$

그리고

$$ \begin{cases} \lambda_{D} \sigma (\beta (r_{\theta} (x,y) - z_{0})) & if & y \sim y_{desirable|x} \\ \lambda_{U} \sigma (\beta ( z_{0} - r_{\theta} (x,y))) & if & y \sim y_{undesirable|x} \end{cases} $$

더 안정적인 학습을 위해, 본 연구에서는 $z_{0}$를 통해 역전파(backpropagation)을 하지 않는다. 이는 순전히 손실의 포화를 제어하기 위해서이다.

직관적으로 KTO는 다음과 같이 작동한다.: 만약 모델이 바람직한 예제의 보상을 무작정 증가시키면, KL 페널티도 증가하여 아무런 진전이 이루어지지 않는다. 이는 모델이 출력을 바람직하게 만드는 요소를 정확하게 학습하도록 강제하는 것이며, KL 항을 평평하게 유지하거나 감소시키면서 보상을 증가시킬 수 있도록 한다. 반대 관점에서도 이 논리는 적용되며, KL 항이 음수가 될 수 없기 때문에 손실 영역에서는 더 빠르게 포화에 도달할 수 있다.

 

KL Estimate

실제로, 위에서 정의한 대로 $z_{0}$를 추정하는 것은 비현실 적이다. 왜냐하면, $\pi_{theta}$에서 샘플링하는 것은 느리고, 인간은 $\pi_{\theta}$에 의해 유도된 전체 분포를 인식하지 않기 때문이다. 인간의 실제 신념에 대한 기대 보상을 계산하면 KL 항에 대한 편향된 추정치를 얻게 되는데, 이는 인간이 최근에 피드백을 준 출력에 과도한 가중치를 부여하는 "이용 가능성 휴리스틱(availability heuristic)"을 가지고 있기 때문이다. 이러한 출력이 $x$의 good continuation일지 여부와는 상관이 없다.

인간이 인식하는 reference point를 더욱 잘 시뮬레이션하기 위해, 우리는 $y_{j}$가 $x_{i}$와 같은 크기의 오프라인 데이터 배치에 속하지만 $ i \neq j $인 $m$개의 쌍 $(x_{i}, y_{j})$를 생성하고, 동일한 배치에 있는 모든 예제에 대해 공통 reference point $z_{0}$를 다음과 같이 추정한다.

$$ \hat{z}_{0} = \max \left( 0, {1 \over m} \sum_{i \neq j} \log {\pi_{\theta}(y_{j} | x_{i}) \over \pi_{ref} (y_{j} | x_{i}) } \right)  $$

이 추정치는 편향되어 있지만, 이는 인간이 자신의 주관적 reference point를 구헝하는 방식에 더 가까워졌기 때문에, 더 바람직한 식이라고 할 수 있다. 또한 각 $x$에 대해 여러 $y^{'}$를 샘플링할 필요가 없기 때문에 효율적으로 계산할 수 있다.

DATA

만약 alignment data가 원래 이진 형태라고 하면, 모든 긍정적 예제는 $y_{desirable|x}$에서 추출되었다고 가정할 수 있고, 모든 부정적 예제는 $y_{undesirable|x}$에서 추출되었다고 가정할 수 있다. 그러나 학술 연구의 standard feedback dataset(HH, SHP, OASST)은 preference 형식으로 되어 있다. 이는 지금까지 가장 효과적이었던 방법들이 선호 기반이었기 때문이다. 따라서 본 연구에서는 선호 데이터를 $y_{w} > y_{l}$로 변환하여, $y_{w}$가 desirable 분포에서, $y_{l}$가 undesirable 분포에서 추출되었다고 가정한다.
이는 단순함을 위한 가정일 뿐이며, preferenc를 이진 feedback으로 더 복잡하게 분해한다면 더 나은 결과를 얻을 수 있지만, 이는 향후 연구로 남겼다. KTO가 선호 기반이 아닌 데이터와도 사용될 수 있음을 보여주기 위해, 일부 실험에서는 정확히 하나의 $y$만을 $x$에 대해 sub-sampling하여 one-y one-x 쌍으로 된 preference 데이터를 제거하고, 데이터 양을 줄이는 방향으로 실험을 수행했다.

 

만약 human feedback이 scores나 ratings 형태로 주어지면, 이를 KTO에 통합하는 가장 간단한 방법은 높은 점수 데이터를 더 많이 가중치로 반영하고, 특정 임계값 이상(이하)의 점수를 받은 예시가 desirable(undesirable)하다고 간주하는 가중치 함수를 구성하는 것이다. 또한 처음부터 점수 기반 HALO를 설계하는 것도 가능하지만, 그러한 손실 함수 설계는 향후 연구과제로 남겨둔다.

Hyperparameter

기본 가중치 함수는 두 개의 hyperparameter $\lambda_{D}$, $\lambda_{U}$를 사용하여 손실 회피 정도를 제어하며, 이 둘은 모두 기본값 1로 설정된다. class imbalance 상황에서, $n_{D}$와 $n_{U}$가 각각 desirable 예시와 undesirable예시의 수를 나타낸다면, 일반적으로 $\lambda_{D}$, $\lambda_{U}$를 다음과 같이 설정하는 것이 가장 좋다.

$$ {\lambda_{D} n_{D} \over \lambda_{U} n_{U}} \in \left[ 1, {4 \over 3} \right] \tag{9}$$

이때, 두 값중 하나는 1로 설정되고, 비율은 다른 값을 조정하여 제어된다. 예를 들어, desirable 예시와 undesirable 예시의 비율이 1:10이라면, $\lambda_{U}=1$로 설정하고, $\lambda_{D} \in [10, 10.33]$ 범위로 설정한다. $[ 1, 4/3 ]$의 구간은 실험적으로 결정했다.

독성 예방과 같이 부정적인 결과를 최소화하는 것이 더 중요한 작업에서는 $\lambda_{D}$, $\lambda_{U}$를 $\lambda_{D}n_{D} < \lambda_{U}n_{U}$가 되도록 설정하는 것이 더 효과적일 수 있으며, 이는 긍정적인 KTO예시가 덜 독성 있는 출력이지만 여전히 피해야 하는 경우에 도움이 되는 걸 알 수 있다. 이는 preference dataset을 단순히 나누어 KTO 데이터를 생성하는 경우에 해당될 수 있다.

또한 $\lambda_{D}$와 $\lambda_{U}$는 중요한 가중치로도 사용할 수 있으며, 이는 offline Dataset $D$가 원하는 feedback 분포와 일치하지 않을 때 유용하다. 기본적으로 실험에서는 $\beta = 0.1$로 설정하며, 손실이 더 빠르게 포화되기를 원한다면, $\beta$를 더 높게 설정하면 된다.

 

현재 설정에서는 $\{ \lambda_{D}, \lambda_{U} \}$를 통해 desirable하거나 undesirable한 예시에 대해 서로 다른 민감도를 허용하며, 이는 수렴시 모든 undesirable 출력이 부정적인 보상을 받고, 모든 desirable 출력이 긍정적인 보상을 받는다는 전제에 기반한다.

이 경우, $\lambda_{U}$와 $\lambda_{D}$는 원래의 Kahneman-Tversky 가치 함수에서 $\lambda$와 직접적으로 대응된다. 그러나, 훈련 중에는 desirable 출력이 부정적인 보상을 받을 수 있고, 그 반대의 경우도 발생 가능하다. 이때, $r_{\theta} (x,y) - z_{0} $가 긍정이던 부정이던 $\lambda_{D}$는 동일하게 적용된다. 본 연구에서는 수렴 시 비대칭성이 가능할 뿐 아니라, 동일한 출력 클래스 내에서도 훈련 중에 비대칭성을 원할 수 있다. 이 경우, $r_{\theta}(x,y) -z_{0}$가 부정적일 때 더 높은 $\lambda_{D}$를 사용하고, 긍정적일 때는 낮은 $\lambda_{D}$를 사용할 수 있다. 또한 데이터에 따라 동적으로 $\beta$를 조정하여 위험 회피 성향을 반영할 수도 있을 것이다. 이러한 창의적이고 동적인 hyperparameter는 향후 연구 과제로 남긴다.

Experiments

실험은 두 가지 평가를 수행했다. 위의 Experiments Setup(3.3)을 따라 (1) GPT-4-0613을 사용하여, test 입력에 대해 aligned model의 생성결과를 SFT의 target output의 sequence와 win-rate를 비교하는 실험; (2) MMLU(0-shot), GSM8K(8-shot CoT), HumanEval(0-shot), BigBench-Hard(3-shot CoT)등의 generation benchmark.

실험 구성은 모든 방법에서 동일했으며, 단 데이터 불균형 실험에서는 KTO에 대해 $\lambda_{D}$와 $\lambda_{U}$를 (9)번 방정식에 따라 설정했다.

KTO $ \geq $ DPO

위의 Figure 3에서 볼 수 있듯, SFT+KTO는 SFT+DPO와 경쟁할만한 성과를 보였으며, Llama-{7B, 13B, 30B} 모델에서는 KTO 단독 성능이 DPO 단독 성능 보다 우수했다. pythia 모델에서는 두 방법 간의 유의미한 차이가 없었는데, 이 차이를 보이기 위해선 최소한의 모델 용량이 필요함을 시사한다.

KTO는 생성 benchmark에서도 DPO 및 다른 baseline보다 더 나은 성과를 보였고, 수학적 추론 데이터셋 GSM8K에서 더욱 성능 향상이 두드러졌다.

At sufficient scale, KTO does not need SFT.

KTO로 align 된 Llama-{13B, 30B} 모델은 SFT+KTO 모델과 비교할 만한 성능을 보였으며, SFT를 먼저 수행하지 않았음에도 불구하고 유사한 성과를 냈다. 본 연구에서 test한 정렬 방법 중에서 이런 결과를 보인 것은 KTO가 유일하다. 이는 KTO만으로도 평균 응답 길이를 대체로 일정하게 유지하는 반면에, SFT 없이 DPO를 실행하면, 응답 길이가 크게 증가할 수 있기 때문이다.

KTO data need not come from preferences.


KTO는 자연적으로 pairs를 이루고 있지 않은 데이터 분포 대신 n개의 preference pairs에서 나온 2n개의 예시를 몰래 활용하고 있을까?

이를 test하기 위해, Llama-7B 모델을 KTO에 align 하기 이전에, 점점 더 많은 desirable data를 무작위로 제거한다. 예를 들어, 원치 않는 데이터는 그대로 두고 desirable data의 90%를 폐기하면, desirable ratio가 1:1에서 1:10이 되며, 비선호 데이터 예시에는 더이상 선호하는 상대가 없어진다. 본 연구에서는 위의 Hyperparameter에서 언급한 기준을 충족하기 위해, $\lambda_{D}=13.33, \lambda_{U}=1$로 변경하여 이러한 불균형을 처리하고, desirable 데이터의 90%를 제거한다. 
위의 그림에서 Llama-7B인 경우, desirable data가 최대 90%까지 실제로 폐기될 수 있으며, DPO를 능가할 수 있음을 발견했다.

 

또한 이 주장을 더 검증하기 위해서, Mistral-7B 모델을 OpenAssistant 데이터셋에 대해 DPO(n pairs 사용), standard KTO(모든 2n개의 Outputs 사용), 그리고 각 $x$에 대해 하나의 $y$만 사용하는 KTO로 align 했다. OpenAssistant에서 하나의 $y$출력이 같은 $x$에 대한 다른 $y$에 의존하지 않기 때문에, 후자는 데이터가 본질적으로 pairs가 없는 분포에서 나온 상황을 반영한다. 각 $x$당 하나의 $y$만을 사용하는 설정은 훈련 데이터 양을 72% 줄였음에도 불구하고, KTO로 정렬된 모델은 여전히 DPO 모델과 official instruction-tuned Mistral-7B보다 더 나은 성능을 보였다.

Changing the design of KTO makes it significantly worse.

KTO의 설계를 변경하면 성능이 크게 저하되는데, 우선 KTO가 HALO로서 자격을 갖추기 위해 필요한 reference point $z_{0}$를 제거하면, BBH와 GSM8K에서 각각 3.6점과 4.0점의 성능 저하가 발생한다. (위의 Table 1. 중간).

KTO가 HALO로 남도록 하는 변경조차도 대부분 최적이 아니다. 예를 들어, 가치 함수의 대칭성을 제거하고 $ 1- \sigma (\cdot)$에서 $- \log \sigma( \cdot)$으로 변경하는 것(즉, DPO처럼 가치 함수를 모든 곳에서 오목하게 만드는 것)은 BBH와 GSM8K에서 각각 9.4점과 11.0점의 성능 저하를 초래했다. 가치 함수를 항등 함수로 설정하여 risk-neutral(위험 중립적)으로 만들면, BBH 성능이 완전히 붕괴된다. 위험 회피 및 손실 회피 hyperparameter를 통해 곡률과 기울기를 조정하면, 작업에 따라 성능이 약간 개선되거나 저하될 수 있지만, $\lambda_{D}$와 $\beta$의 극단적인 선택은 좋은 성능을 가져올 수 없다.

KTO works without a reference model or SFT, but not as well as standard KTO.

KTO는 참조 모델이나 SFT없이 작동하지만, standard KTO만큼 잘 작동하지 않는다. reference model을 메모리에 저장하지 않으면 훨씬 더 메모리 효율적인 align이 가능하다. 이를 단순하게 구현하는 방법은 $\pi_{ref}$가 모든 $x$에 대한 output의 균등 분포를 반환한다고 가정하는 것이다. 이렇게 하면, $r_{\theta} - z_{0}$는 $ \log \pi_{\theta} (y|x) - H( \pi_{\theta} (y^{'} |x ))$ ($H$는 entropy로 정의하고, $\max \left( 0, {1 \over m} \sum_{i \neq j} \log \pi_{\theta} (y_{j} | x_{i}) \right)$로 추정한다.)

Table 1 중간에서 $\lambda_{D}=1.75$로 설정하면 메모리 효율적인 KTO 변형이 일부 작업에서는 DPO보다 더 나은 성능을 보이고 다른 작업에서는 더 나쁜 성능을 보였지만, 여전히 standard KTO 성능엔 미치지 못했다. 또한 이 방법은 손실 회피 hyperparameter에 더 민감한 것으로 나타났다. $\lambda_{D}$를 {1.5, 2.0}으로 설정하면, GSM8K, BBH에서 성능이 떨어졌다. 그럼에도 불구하고, 이 방법은 최근 제안된 reference-free 방법인 ORPO를 능가했으며, $\pi_{ref}$를 메모리에 로드할 필요가 없고 KTO의 outputs의 batch가 ORPO/DPO에서 사용되는 2m outputs 배치보다 작기 때문에 기존의 모든 접근 방식보다 메모리를 덜 사용한다.

Theoretical Analysis

KTO는 이진 Feedback이 더 약하더라도, preference에 비해 훨씬 더 풍부하고, 저렴하며, 빠르게 수집할 수 있는 데이터를 대량으로 사용하여 이를 보완할 수 있다는 동기로 설계됐다.

이 KTO는 동일한 preference data를 사용한 DPO보다 동등하거나 더 나은 성능을 보이는 이유는 데이터 효율성 뿐만이 아니다.

$\rightarrow$ 각 $x$당 하나의 $y$만 사용하는 설정에서 KTO가 여전히 우수한 성능을 보임

이 현상에 대한 두 가지 이론적 설명은 다음과 같다.

  1. 선호도 likelihood를 최대화하는 것이 반드시 인간의 효용을 최대화하지 않는다.
  2. KTO는 실제 피드백에서 흔히 발생하는 noise & intransitived data를 무시한다.

Proposition 4.1. 현재 policy가 암시하는 보상이 $ \pm \infty$로 수렴할 때, $\pi_{\theta}$의 KTO Update는 0으로 수렴한다.

 

이는 현재 policy $\pi_{\theta}$에 의해 $(x,y)$가 학습하기에 너무 어렵거나 쉽다고 판단될 때, 해당 데이터가 사실상 무시된다는 것을 의미한다. real-world의 Feedback은 매우 noisy하기 때문에, desirable 예시가 매우 부정적인 reward를 출력하는 이유는 잘못된 labeling 때문일 수 있다. KTO는 이러한 학습하기 어려운 데이터를 피함으로써, noise에 맞추는 것을 방지한다. 그러나 이는 KTO가 학습하기 어렵지만 $r^{*}$를 복구하는데 필요한 데이터를 무시하게 되어 복잡한 분포에 대한 과소적합이 발생할 수 있다는 것을 의미하기도 한다. 이러한 과소적합은 더 낮은 $\beta$ 값으로 모델을 조정하고 epoch 수를 높이는 방법으로 완화시킬 수 있다.

 

Proposition 4.2. 가치 함수가 로지스틱 함수라고 가정했을때, $\mathbb{E}_{x \in D, y \in \pi_{\theta}} [r_{\phi}(x,y)] - \beta D_{KL}(\pi_{\theta} (y|x) || \pi_{ref}(y|x))$를 최대화하는 reward function $r_{a}^{*}$에 대해, 동일한 최적의 policy $\pi^{*}$와 동일한 Bradely-Terry preference 분포를 유도하지만, 다른 인간 가치 분포를 가지는 동일한 class의 보상함수 $r^{*}_{b}(x,y) = r^{*}_{a}(x,y)+h(x)$가 존재한다. (여기서 $h(x)$는 어떠한 함수이다.)

 

Rafailov et al. (2023)의 주요 통찰 중의 하나는 동일한 동등 class에 속하는 reward functions (즉, 오직 입력에 따라 달라지는 input-specific component) $\mathbb{E}_{x \in D, y \in \pi_{\theta}} [r_{\phi}(x,y)] - \beta D_{KL}(\pi_{\theta} (y|x) || \pi_{ref}(y|x))$ 이 식 아래에서 동일한 최적 policy와 동일한 Bradely-Terry preference 분포를 유도한다는 것이다. 그러나 본연구는 mild한 가정 하에서 이러한 입력별 변화가 가치 분포 즉, 인간의 효용에 영향을 미친다는 것을 보여준다. 따라서 preference likelihood를 최대화 하는 것이 인간의 효용을 최대화하는 것은 아니다 라는 결론을 내린다. 이것은 KTO와 DPO 간의 격차가 automated LLM-as-a-judge evaluations에서 더 크게 나타나는 이유를 설명하는데 도움이 된다.(Appendix D)

 

Proposition 4.3. outputs $\{ y_{a}, y_{b} \}$인 input $x$에 대해, Dataset $D$는 상충하는 preferences $y_{a} > y_{b}$와 $y_{b} > y_{a}$를 각각 비율 $p \in (0.5, 1)$ 및 $(1-p) \in (0,0.5)$로 비율화 한다고 하자. 만약 $p^{1/ \beta} \pi_{ref} (y_{b} | x)$라면, 최적 DPO policy는 minority-preferred y_{b}를 더 많이 생성할 가능이 있으며, 최적 KTO policy는 loss-neutral value function $( \lambda_{D} = \lambda_{U} )$에서 majority-preferred $y_{a}$만을 strictly 생성할 것이다.

 

비 공식적으로, $x$에 대한 output에 대해 모순되는 두 개의 preference가 있고, 다수는 비율 $p$로 $y_{a}$를 선호하고, 소수는 $1-p$로 $y_{b}$를 선호한다고 하자. 최악의 경우, $p$가 충분히 낮고 reference model이 충분히 align되지 않았을 때, 최적 DPO policy는 암시된 보상 $r_{\theta} (x, y_{b}) > r_{\theta} (x, y_{a}$ 임에도 불구하고 minority-preferred output $y_{b}$를 더 많이 생성할 가능성이 있다. 반면, optimal KTO policy는 가치 함수가 loss-neutral $(\lambda_{D} = \lambda_{U})$일 경우, 즉, 기본 설정에서 majority-preferred $y_{a}$를 생성할 것이다. 이는 KTO가 피드백의 비일관성을 처리할 때, 최악의 경우, 더 나은 결과를 제공한다는 것을 시사한다.

KTO vs. DPO - when to use which?

인간 Feedback이 binary 형식일 때, 특히 desirable 예시와 undesirable 예시의 불균형이 있을 경우, KTO를 선택하는 것이 맞다. 하지만, 데이터가 preferences 형태 일 때는 명확한 선택이 어렵다. preference의 noise와 비일관성이 충분히 작을 경우에는 DPO가 더 나은 성능을 발휘할 것이다. 이는 KTO가 과소적합할 위험의 이유가 있기 때문이다. 이 위험은 $\beta$ 값을 낮추고 더 많은 epoch 동안 KTO를 align 함으로써 완화할 수 있다. 하지만 noisy와 비일관성이 충분히 존재한다면, KTO의 worst-case guarantees 이 더 나을 것이다. 대부분의 공개된 preference dataset(예: SHP, OpenAssistant)은 여러 사람의 feedback에서 비롯된 noise 가 많고 이들의 preference 가 어느 정도 모순(상충)될 가능성이 있다. 이것이 KTO가 DPO의 성능을 따라잡거나 능가할 수 있었던 이유이다. 심지어 합성된 feedback조차도 noise와 비일관성을 포함할 수 있으며, 이것이 UltraFeedback과 align할 때, KTO가 DPO를 능가할 수 있는 이유이다.