Setting

통계학적인 관점에서, 강화학습(RL)은 uncertainty 하에서 이루어지는 sequential decision making 혹은 dynamic optimization 이라고 할 수 있다. 일반적인 supervised learning 세팅에서는 관측 데이터 $(X,Y)$ 쌍들로부터 conditional distribution $\Pr(Y|X)$ 를 학습하는 것인 반면, RL에서는 보상(reward)을 최대화하는 policy $\pi$를 찾고자 한다.

Policy

Parametric Conditional Distribution

Policy는 상태 $s$에서 행동 $a$를 선택할 확률을 나타내는 조건부 확률 분포이다.

• Input: state $s\in \mathcal{S}$, 현재 상태(ex. 현재 게임 화면, 채팅 기록 등)
• Output: action $a \in \mathcal{A}$, 취할 수 있는 행동(ex. 캐릭터 이동, 답변 생성에서 Token "The" 선택 등)
• Parameter $\theta$ : weights of neural networks

Reward

Reward function $r_t(s)$는 시점 $t$에서 상태가 $s$인 agent가 환경으로부터 받는 scalar feedback signal이다. 강화학습의 목표는 장기적으로 받을 보상의 합을 최대화하는 것이다.

• $\gamma$ : discount factor, $0 \leq \gamma < 1$ (주로 0.99 사용), 미래 보상에 대한 현재 가치의 감소율

Value function

다음과 같은 3개의 주요 가치 함수(value function)가 있다.

• State-Action Value Function $Q^{\pi}(s,a)$: 상태 $s$에서 행동 $a$를 취했을 때 기대되는 누적 보상
• Value Function $V^{\pi}(s)$: 상태 $s$에서 기대되는 누적 보상 (평균적인 action 기준)
• Advantage Function $A^{\pi}(s,a)$: 특정 행동의 가치가 평균적인 행동에 비해 얼마나 좋은지

Optimization Objective

강화학습의 목표는 policy $\pi_{\theta}$의 파라미터 $\theta$를 최적화하여 기대 누적 보상(expected cumulative reward)을 최대화하는 것이다. 즉, 다음과 같은 목적함수를 최적화한다.

PPO(Proximal Policy Optimization)

PPO는 제약 조건 하에서의 정책 최적화 알고리즘으로 볼 수 있다.

Vanilla Policy Gradient

목적함수 $J(\theta)$의 gradient estimator는 다음과 같이 표현된다.

$G_t$ 가 아닌 $A_t$ 를 사용하는 이유는 $G_t$가 가지는 높은 분산(variance)을 줄이기 위함이다. (highly stochastic reward 환경에서 특히 유용)

Trust Region Policy Optimization

TRPO는 policy 업데이트 시 KL divergence 제약 조건을 도입하여 급격한 변화를 방지한다.

이때 위 제약조건 문제는 penalized objective로 변환하여 해결할 수 있다.

Clipped Surrogate Objective

PPO는 TRPO의 복잡성을 줄이기 위해 clipped surrogate objective를 도입한다. TRPO의 objective 내 비율 항목을 다음과 같이 정의한다.

그러면, CLIP된 목적함수는 다음과 같이 표현된다.

여기서 $\epsilon$는 작은 양수 하이퍼파라미터로, 일반적으로 0.2로 설정된다.

만일 action이 advantage $\hat{A}_t$에 대해 긍정적이라면, 확률을 높인다 ($r_t(\theta) > 1$). 다만, $r_t(\theta)$가 $1 + \epsilon$를 초과하지 않도록 클리핑한다.

• Schulman, John, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. 2015. “Trust Region Policy Optimization.” Proceedings of the 32nd International Conference on Machine Learning, June 1, 1889–97. https://proceedings.mlr.press/v37/schulman15.html.

• Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017. “Proximal Policy Optimization Algorithms.” arXiv:1707.06347. Preprint, arXiv, August 28. https://doi.org/10.48550/arXiv.1707.06347.

LLM에서의 Policy Optimization

위에서 정의한 policy optimization 관점은 LLM alignment에도 거의 그대로 대응된다. 차이는 환경(environment) 과 행동(action) 의 해석이다. 일반적인 게임 RL에서는 상태가 게임 화면이고 행동이 조작 입력이라면, LLM에서는 상태가 prompt와 지금까지 생성한 context이고 행동은 다음 token 선택이다.

따라서 LLM의 policy optimization은 다음과 같이 볼 수 있다.

• Policy: language model $\pi_\theta(y|x)$
• State: prompt $x$와 partial completion $y_{<t}$
• Action: next token $y_t$
• Reward: 사람 선호, 안전성, 유용성 등을 반영한 scalar score $r(x,y)$
• Reference policy: 보통 SFT가 끝난 모델 $\pi_{ref}$ 혹은 이전 policy $\pi_{old}$

여기서 중요한 점은 reward가 token마다 dense하게 주어지는 경우보다, 완성된 답변 $y$ 전체에 대해 하나의 점수로 주어지는 경우가 많다는 것이다. 즉, LLM alignment에서는 episodic RL에 가까운 세팅으로,

와 같은 목적함수를 생각할 수 있다. 하지만 reward만 최대화하면 모델이 원래의 language modeling 능력에서 크게 벗어나거나 reward model을 exploit할 수 있으므로, 실제로는 reference model과의 KL divergence를 함께 제어한다.

직관적으로는 보상은 높이고 싶지만, SFT 모델에서 너무 멀어지지는 않게 하는 것이다.

RLHF Pipeline: SFT $\rightarrow$ Reward Model $\rightarrow$ PPO

InstructGPT류의 RLHF pipeline은 대략 다음 3단계로 정리할 수 있다.

1. Supervised Fine-Tuning (SFT)
   • 사람이 작성한 prompt-response pair로 language model을 supervised learning한다.
   • 이 모델이 이후의 reference policy $\pi_{ref}$가 된다.
2. Reward Model (RM) 학습
   • 같은 prompt에 대해 여러 답변을 만들고, 사람이 선호 순위를 매긴다.
   • reward model $r_\phi(x,y)$는 어떤 답변이 더 선호되는지 예측하도록 학습된다.
3. PPO로 policy 최적화
   • 현재 policy $\pi_\theta$가 답변을 생성한다.
   • reward model이 점수 $r_\phi(x,y)$를 준다.
   • KL penalty를 포함한 reward로 PPO 업데이트를 수행한다.

Reward model은 보통 pairwise preference data로 학습한다. prompt $x$에 대해 사람이 선호한 답변을 $y_w$, 덜 선호한 답변을 $y_l$라고 하면 Bradley-Terry 형태의 loss를 사용할 수 있다.

이후 PPO 단계에서는 reward model score에 KL penalty를 더한 shaped reward를 쓴다.

Token-level로 쓰면 다음과 같이 각 token의 log-ratio를 합산한 항이 KL penalty 역할을 한다.

결국 LLM에서 PPO는 다음과 같은 형태의 optimization으로 이해할 수 있다.

단, 여기서 $r_t(\theta)$는 token action에 대한 확률비이다.

LLM PPO에서 value function은 prompt와 partial completion을 보고 남은 return을 예측하는 critic으로 붙는다. 즉, actor는 next-token distribution을 내고, critic은 $V_\psi(x,y_{<t})$를 예측한다. Advantage는 reward model score와 KL penalty를 반영한 return으로부터 GAE(Generalized Advantage Estimation) 등을 통해 추정한다.

왜 LLM에서 PPO가 어려운가?

PPO는 이론적으로는 자연스럽지만 LLM에 적용할 때 꽤 무겁다.

• Sampling cost: policy가 직접 답변을 생성해야 하므로 rollout 비용이 크다.
• Reward model dependency: reward model이 잘못 학습되면 reward hacking이 발생할 수 있다.
• Actor-critic complexity: policy model, reference model, reward model, value model을 함께 다뤄야 한다.
• Hyperparameter sensitivity: KL coefficient $\beta$, clipping range $\epsilon$, learning rate, advantage normalization 등에 민감하다.
• On-policy 성격: 새 policy에서 생성한 sample로 업데이트해야 하므로 supervised fine-tuning보다 data efficiency가 낮다.

이러한 이유로, PPO는 RLHF의 대표적인 방법이지만 구현과 운영 비용이 큰 편이다. 이 지점에서 등장하는 더 단순한 대안 중 하나가 DPO이다.

DPO(Direct Preference Optimization)

DPO의 핵심 아이디어는 reward model을 명시적으로 학습한 뒤 PPO로 최적화하지 말고, preference data로부터 policy를 직접 학습하자는 것이다. 즉, RLHF pipeline의 RM 학습과 RL optimization 단계를 하나의 supervised-looking loss로 합친다.

출발점은 KL-regularized reward maximization이다.

이 문제의 최적 policy는 다음과 같은 Boltzmann distribution 형태를 가진다.

이를 reward에 대해 정리하면,

가 된다. 같은 prompt $x$에서 두 답변 $y_w,y_l$를 비교하면 partition function $Z(x)$는 서로 상쇄된다.

DPO는 여기서 최적 policy $\pi^*$를 현재 학습 중인 policy $\pi_\theta$로 parameterize하고, Bradley-Terry preference likelihood에 대입한다. 따라서 DPO loss는 다음과 같다.

Sequence log-probability는 token log-probability의 합이다.

이 loss의 의미는 단순하다.

• 선호 답변 $y_w$에 대해서는 $\pi_\theta(y_w|x)$를 $\pi_{ref}(y_w|x)$ 대비 더 크게 만든다.
• 비선호 답변 $y_l$에 대해서는 $\pi_\theta(y_l|x)$를 $\pi_{ref}(y_l|x)$ 대비 더 작게 만든다.
• 단, $\beta$가 reference policy에서 얼마나 멀어질지 조절한다.

DPO vs PPO

요약하면 PPO는 reward model이 정의한 보상을 직접 최대화하는 RL 방법이고, DPO는 선호 비교 데이터를 이용해 KL-regularized optimal policy를 직접 fitting하는 방법이다. PPO가 더 일반적인 RL framework에 가깝다면, DPO는 LLM preference tuning에 특화된 더 단순한 objective라고 볼 수 있다.

한 줄 직관

• PPO: "이 답변의 reward가 높으니, 그 답변을 만들 확률을 조금 높이자. 단, 너무 급격히 바꾸지는 말자."
• DPO: "같은 prompt에서 사람이 고른 답변은 reference 대비 더 가능하게, 버린 답변은 reference 대비 덜 가능하게 만들자."

마무리

Policy optimization은 강화학습의 일반적인 수학적 틀이지만, LLM에서는 next-token prediction model을 policy로 해석하면서 alignment 방법론과 직접 연결된다. PPO는 RLHF의 고전적인 접근으로 reward model과 KL-regularized policy update를 사용하고, DPO는 preference likelihood를 통해 그 과정을 더 직접적인 supervised objective로 바꾼다. 결국 두 방법 모두 핵심은 동일하다. 사람이 선호하는 응답의 확률은 높이고, 원래 모델의 언어 능력과 분포에서는 너무 멀어지지 않도록 제어하는 것이다.

• Ouyang, Long, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, et al. 2022. “Training Language Models to Follow Instructions with Human Feedback.” arXiv:2203.02155. Preprint, arXiv, March 4. https://doi.org/10.48550/arXiv.2203.02155.

• Rafailov, Rafael, Archit Sharma, Eric Mitchell, Christopher D. Manning, Stefano Ermon, and Chelsea Finn. 2023. “Direct Preference Optimization: Your Language Model Is Secretly a Reward Model.” Advances in Neural Information Processing Systems 36. https://arxiv.org/abs/2305.18290.

• Christiano, Paul F., Jan Leike, Tom B. Brown, Miljan Martic, Shane Legg, and Dario Amodei. 2017. “Deep Reinforcement Learning from Human Preferences.” Advances in Neural Information Processing Systems 30. https://arxiv.org/abs/1706.03741.