Why DPO is a Misspecified Estimator and How to Fix It 논문 리뷰

Summary

DPO(Direct Preference Optimization)은 전통적인 two-stage RLHF에서 선호 데이터를 이용하여 보상모델을 학습하고, 그 보상을 최대화하도록 정책을 업데이트하는 방식과 다르게 선호쌍을 분류하는 형태의 supervised loss로 정책을 바로 업데이트한다. 구현이 매우 단순하며 보상 모델에 의존적이지 않고, 특히 LLM에서 그 효과를 인정받아 널리 쓰이는 방법론이다.

 

그러나 본 논문은 DPO는 LLM처럼 parameterized된 정책 클래스에서는 실제 보상을 타깃하여 훈련되지만 정책이 표현 가능한 보상 공간으로 사영(project)할 뿐임을 보이며, 이로 인해 선호 순서가 뒤집히거나 정책의 trajectory가 만들어내는 평균적인 보상이 감소하는 결과가 일어난다는 것을 밝힌다.

 

따라서 본 논문의 저자들은 기하학적 분석을 바탕으로 AuxDPO를 제안하여, 보상 공간에 추가 자유도를 부여하는 방식으로 misspecification을 완화하고자 한다.

Preliminaries

Direct Preference Optimization

KL regularized RL objective는 다음과 같으며,

해당 objective를 최적화하는 policy를 closed-form으로 구할 수 있으며, 해당 closed-form으로 reward function에 관한 식으로 도출할 수 있다. 여기서 핵심은 보상이 주어졌다면, 그 보상은 reference policy와 target policy의 우도비임을 알 수 있다.

여기서 주목할 점은 reward function은 더는 고정된 함수가 아니라는 점이다. Reward function은, reference model이 고정되어 있는 경우 target policy가 내난 확률값에 따라 정의된다. 그러면 선호 쌍 데이터를 바탕으로, win output을 lose output보다 더 선호할 확률은 다음과 같은 BTL 모델로 표현 가능하다:

논문에서 언급하는 구절에 대해서 살펴보자(이 논문을 이해할 주요 구절이다):

DPO에서 제시하는 empirical loss는 BTL model의 negative log-likelihood를 최소화하는 식이다. 그러나 "만약에" 실제 보상이 주어져있다고 가정하면, 사실은 실제 보상에 대해서 모델이 추정하는 보상이 align 되도록 하는 cross-entropy loss가 되는 것이다. 여기서

부분에 집중해보자. Reward function의 output은 scalar이라 1-dim이지만, 여기서는 정책이 표현 가능한 reward의 차원을 의미한 것이라고 보면 된다. 여기서 논문이 주로 다루는 misspecification이 나오는 것인데, 그 이유는 policy가 실제로 간접적으로 표현할 수 있는 reward의 차원에 비해 state, action의 차원이 압도적으로 크기 때문이라고 한다. 자세한 이유를 다음 파트에서 알아보도록 하자.

 

Reward Misspecification in DPO

다음 Proposition을 보자:

(4)번 수식의 의미는, 진짜 보상 $r^*$가 만들어내는 BTL 비교 확률과 policy가 간접적으로 만드는 BTL 비교 확률 사이의 KL divergence를 최소화한다는 것인데, 여기서 핵심은 KL이 실제 보상 $r^*$의 state, action space에서 이루어지는 것이 아닌, 정책이 표현하게 되는 state, action space 집합 안에서만 계산된다는 것이다. 논문이 이것이 projection이라고 말하는데, 이러한 이유는 고정된 reward를 제약된 집합 위로 KL 기준으로 사영하기 때문이다.

 

만약 여기서 실제 보상의 차원과 모델의 표현 차원이 같다면 projection은 자기 자신이 되어 RLHF에서 제공하는 closed-form 정책과 일치하게 되어 정확하지만, 대부분은 모델의 표현 차원에 비해 실제 보상의 차원이 압도적으로 크기 때문에 모델이 진짜 보상을 표현할 수 없는 misspecified estimation setting에 있게 되는 것이다.

 

이로 인한 문제는 두 가지로 요약해볼 수 있다. Pairwise KL에서, 데이터 분포에 따라 분포가 큰 곳의 alignment는 잘 이루어지지만, 그렇지 않은 곳의 경우 alignment가 잘 이루어지지 않거나, 심지어 특정 비교 순서가 뒤집히는 상황도 발생할 수 있다. 특히 KTO 논문에서 실험적으로 데이터의 비율에 따라 DPO의 성능이 달라지는 것을 확인한 바가 있다.

또한 평균 보상도 줄어들 수 있다. 고전적 two-staged RLHF는 policy improvement가 보장되었지만, DPO는 엄밀히 말해서 reward를 최적화하는 것이 아닌 KL을 최적화하는 것이기 때문에, 비교 확률을 맞추는 것과는 별개로 실제 기대 보상을 줄일 수 있는 것이다.

 

본 리뷰에서는 3.1 LOCAL GEOMETRY OF DPO를 다루지 않고 가지만 확인하면 좋을 것이다. 결과만 말하자면 policy의 implicit reward manifold는 local하게 선형 공간이며, DPO는 실제 reward를 그 공간으로 사영하고 그 사영 방향은 데이터의 빈도에 따라 결정된다. 그러한 결과로 순서가 뒤집히거나, 최적 행동 확률이 줄거나 평균 보상이 감소할 수 있으며, 이는 데이터 부족이나 optimizer의 문제도 아닌 policy capacity 제약 자체에서 오는 기하학적 문제인 것이다.

 

Towards Mitigating DPO's Pitfalls

논문에 수식이 많아서 어지러우니, 우선 전체 순서를 요약하고 덤벼보도록 하자.

1. RLHF를 locally 분석
2. reward equivalence class 개념 등장
3. DPO linearization과 RLHF equivalence class의 관계 규명
4. misspecification의 정확한 구조 파악
5. nullspace 방향을 열어주면 해결 가능
6. 그것이 AuxDPO

이것을 단계별로 정의해보자.

 

RLHF objective $J(\theta; r^*)=\mathbb{E}[r^*] - \beta D_{KL}(\pi_\theta || \pi_{\theta_0})$를 Taylor expansion으로 근사해보자. TRPO 공부 좀 했으면 금방 한다. ㅎㅎ

여기서 $F$는 Fisher information matrix이고, $A=\nabla\log\pi_{\theta_0}(a|s)$이다. 결과적으로 다음과 같은 형태가 된다:

이 근사 형태를 잘 풀어주면 natural policy gradient가 다음과 같이 결정됨을 알 수 있다:

이 식에서 중요한 것은, $Ar^*$가 정책의 업데이트를 결정한다는 점이다. 따라서, $r_1, r_2$가 $Ar_1 = Ar_2$이면 동일 정책을 만드는 것이고, 이 때 equvalance class는 $r_1-r_2 \in N(A)$라는 것이다.

Proposition 7은 DPO의 선형화된 reward $r^{\beta}_\theta = \beta A_{\theta_0}^\top (\theta - \theta_0)$는 RLHF의 equivalence class 중 Mahalanobis norm이 최소인 대표 원소를 선택하는 것을 말한다. Misspecification의 본질은 여기서 나오는데 실제 reward는 DPO가 표현 가능한 것 + nullspace 성분으로 구성되어 있다는 것, 즉 $r^*=r^\beta_{\theta^*}+\delta, \delta\inN(A)$라는 것이다. 다시 말해서, DPO는 RLHF reward equivalence class 중 최소-norm 대표만을 표현할 수 있기 때문에, 진짜 reward에 존재하는 nullspace 성분을 모델링하지 못한다. 그 결과 BTL 확률을 맞추는 과정에서 왜곡이 발생한다는 것이다.

 

따라서 본 논문은 nullspace 방향을 열어주는 AuxDPO를 제시한다. 수식을 확인해보자:

기존에 policy에서 implicitly induced reward에, null space에 속하는 $\delta$를 추가하여 loss를 구성한 방식이다. Null space는 기본적으로 찾는 것이 상당히 어렵기 때문에, 논문에서는 penalty term으로 null space를 대체하는 방식을 취한다:

아래의 penalty term은 $A\delta\approx0$이 되도록 강제하는, 즉 policy에 영향을 주지 않는 방향(nullspace)으로만 존재하도록 학습된다. 

Experiments

논문에서 주장한 바와 같이, 실제로 기존 방법들에 비해 유의미하게 높은 성능을 거두었으며 특히 OOD setting에서 우수한 성능을 보이는 것을 볼 수 있다.

 

논문 리뷰는 여기서 마무리하도록 하겠다!