Direct Alignment from Preferences - Part 01. RLHF

19 Feb 2024

Introduction

2023 년은 감히 chatGPT 의 시대였다고 말할 수 있을 것 같습니다. chatGPT 의 성공에는 다양한 요소들이 있습니다. GPU 의 발전과 함께 엄청나게 거대한 모델을 빠르게 학습할 수 있게 되면서 175B, 그 이상의 모델이 등장하게 되었습니다. 또 단순히 많은 데이터를 모으는 것을 넘어 높은 퀄리티의 정제된 Instruction 데이터셋들을 사용하면서 사람과 상호작용할 수 있는 채팅 모델로써의 성능을 끌어올렸습니다. 그러나 많은 요인들 중 단 하나의 요인을 꼽아보라면 저는 RLHF, 즉 Model alignment 를 꼽겠습니다.

Model alignment 는 간단히 말해 “사람의 취향에 맞는 답변을 생성하도록 모델을 학습“하는 것입니다. 기존의 SFT (Supervised Fine-Tuning) 는 모델이 학습데이터와 같은 답변을 생성하도록 학습됩니다. 그러나 SFT로 학습된 모델은 미처 학습하지 못한 데이터에 대해 적절하게 대응하기 어렵습니다. 또한, 학습 데이터에 적절치 못한 답변이 포함된다면 부적절하거나 왜곡된 답변을 생성할 위험이 있습니다. 따라서 Model alignment 는 인간의 피드백을 기반으로 모델의 답변을 개선합니다.

chatGPT 에서 사용된 Model alignment 방식은 RLHF 라는 방식으로 RM (Reward Model) 의 학습과 RL (Reinforcement Learning) 을 이용한 학습방법으로 널리 알려져 있습니다. 하지만 RL 은 매우 불안정한 학습 방법입니다. 따라서 RL 을 사용하지 않고 학습 안정성을 높이는 여러 방법론들이 등장했고, 이들을 DAP (Direct Alignment from Preference) 라고 부릅니다.

이번 포스팅 3부작(아마?)에서는 RL과 DAP 가 어떻게 다른지 알아보고, 최신 RL, 그리고 DAP 방법론들에 대해서 알아보도록 하겠습니다. Part 1 은 RL 방법론들입니다.

Model alignment - RL

RLHF

OpenAI 가 2022년에 내놓은 Training language models to follow instructions with human feedback (Ouyang et al., 2022) 논문에서 제안된 방법론입니다. 사실 그 전에도 Learning to summarize from human feedback (Stiennon et al., 2020) 논문을 통해 Summarization task 에 대해 동일한 방법을 제안했지만, chatGPT 만큼의 파급력을 보이지는 못했습니다.

RLHF 는 3개의 step 으로 나뉩니다. 1) SFT: Chat 형태로 구축된 데이터를 통해 모델을 학습합니다. 2) RM: 하나의 질문에 대해 여러개의 응답을 생성하고 사람은 그 중 어떤 응답을 선호하는지 태깅합니다. 3) RL: 학습된 RM 과 RL 알고리즘 (PPO) 을 이용해 모델을 개선합니다.

이 중 Step 2와 Step 3에 대해 조금 더 디테일하게 수식을 살펴보도록 하겠습니다.

• Reward model training
  • Bradley-Terry Preference model
    • 선호도를 반영하여 학습된 Reward 함수 $r(x,y)$ 가 있을 때 $y_1$이 $y_2$ 보다 선호될 확률은 아래와 같이 표현됩니다. Reward output 의 Softmax 라고 이해하면 될 것 같습니다. 간단하게 말해, 상대편보다 reward 가 높으면 높을 수록 선호될 확률은 높습니다.
    \[p^*(y_1 > y_2 | x) = \frac{\exp(r^*(x, y_1))}{\exp(r^*(x, y_1)) + \exp(r^*(x, y_2))}.\]
  • Reward model loss function
    • 하나의 질문에 대해 두개의 응답쌍으로 구성된 데이터를 만들고, 사람이 선호도를 태깅합니다. 선호된 응답을 $y_w$, 반대 응답을 $y_l$ 로 정의합니다.
    • 아래 Loss 를 사용하여 모델을 학습합니다. $y_w$ reward 는 높아지고 $y_l$의 reward 는 낮아지는 방향으로 모델이 학습됩니다.
    \[\mathcal{L}_R(r_\phi, D) = -\mathbb{E}_{(x,y_w,y_l) \sim D}[\log \sigma(r_\phi (x, y_w) - r_\phi (x, y_l))]\]
• RL (PPO)

  • PPO 알고리즘을 사용해 아래의 Objective function 을 만족하는 모델을 학습합니다.

    \[\max_{\pi_\theta} \mathbb{E}_{x \sim D, y \sim \pi_\theta(y|x)}[r_\phi (x, y)] - \beta D_{KL}[\pi_\theta(y | x) || \pi_{\text{ref}}(y | x)]\]
  • Reward 를 최대화 하면서 기존의 reference 모델과 KL-divergence 가 너무 커지지 않도록 학습합니다.
  • KL-divergence 를 regularizer 로 사용하는 것은 기존에 학습된 지식들을 너무 많이 잃지 않도록 하는데 있습니다. 또한 Regularizer 를 사용하지 않으면 RM 에 over-optimized 되는 reward hacking 현상이 발생할 수 있습니다.

    🚨 LLM 에서 Reward hacking 이란?
        
    “무엇을 도와드리면 될까요?”, “그렇게 한 번 해보겠습니다. 어떻게 하면 좋을지 알려주세요!” 와 같이 쓸데없지만, 거짓말을 하는 답변보다는 선호될 확률이 높은, 사실상 무의미한 답변들을 계속 생성해버리는 현상
    

해당 방법론은 효과적이었지만 3가지 치명적인 단점을 가지고 있습니다.
1) Expensive human annotator

• 사람이 평가한 preference 데이터를 구축하는 것은 비용적으로도, 시간적으로도 굉장히 오래 걸리는 문제입니다.
• 현재 많은 open-source 데이터가 존재하기는 하지만, 고도로 정제된 한국어 preference 데이터는 찾아보기 어렵습니다. 대부분 영어 번역본들입니다. 2) Expensive computational cost
• PPO 알고리즘을 위해서는 4개의 모델이 필요합니다: Policy model, Reference model, Reward model, Value model
• Policy model 만 해도 175B에 육박하는 상황에서 추가로 비슷한 크기의 3개 모델을 더 띄울만한 GPU 공간을 확보해야 합니다.

이 첫 번째 문제를 해결하기 위해 2023년 Google Research 에서는 RLAIF 라는 방법론을 제안합니다.

RLAIF

Google Research 에서 2023년에 발표한 RLAIF: Scaling Reinforcement Learning from Human Feedback with AI Feedback (Lee et al., 2023) 논문에서 사용된 방법론입니다.

컨셉은 매우 간단합니다. RLHF 에서 Human feedback 을 AI feedback 으로 대체하겠다는 것인데요. 이는 Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena (Zheng et al., 2023) 에서도 주장되는 내용으로, LLM 이 충분히 사람의 평가와 유사한 성능을 낼 수 있음을 이용한 것입니다.

chatGPT, Bard, Claude 같은 Off-the-shelve LLM 들에게 평가 프롬프트를 제공해 답변의 선호도를 태깅하는 것이 사람이 직접 태깅하는 것과 높은 correlation 을 가진다

실제로 RLAIF 를 사용할 때와 RLHF 를 사용할 때의 성능 차이는 크게 나지 않았지만

가격 면에서는 AI labeler 가 0.06 $/example 인데 반해 Human labeler는 0.67 $/example로 11배가 넘게 차이난다고 주장하였습니다.

또한 AI feedback 은 가격 뿐만 아니라 속도에도 이점이 있어 Iterative training 을 가능하게 합니다. RLHF 의 경우 학습 사이클이 완료된 후 학습된 모델의 응답들을 모아 human annotator 에게 다시 태깅을 요청하는 과정에서 시간이 많이 소요됩니다. 그러나 RLAIF 는 AI feedback 을 언제든 할 수 있기 때문에 학습 사이클을 몇 번이고 빠르게 수행할 수 있다는 장점이 있습니다.

ReST (Deepmind)

Deepmind 에서는 Reinforced Self-Training (ReST) for Language Modeling (Gulcehre et al., 2023) 논문을 통해 ReST 라는 방법론을 제안했는데요, 이는 RLHF 를 iterative 하게 수행할 수 있는 방법론을 제안했습니다.

ReST 는 4개의 step 으로 구성됩니다. 1) RM training: HF 데이터로 RM 학습 2) Grow step: Current policy model을 이용해 하나의 prompt 에 대한 응답 여러 개 생성 3) Data tagging: RM 을 이용해 preference 태깅

• 이 때 reward threshold 보다 reward가 높은 데이터만 다음 step에 사용 4) Improve step: step 3에서 모은 데이터와 Offline RL 방법론을 사용해 policy 학습 5) 2-4 step 반복

처음 만들어 놓은 RM 을 활용하여 자체 제작한 응답들에 대해서도 preference 태깅을 할 수 있다는 것입니다. 다만 아래 그림과 같이 iteration 을 반복할 수록 reward threshold 를 점점 높여가며 data quality 를 향상시켜야 성능이 점진적으로 향상될 수 있다고 합니다.

또한 여러 번의 iteration 을 돌 때에는 learning rate 를 작게 유지해서 모델이 이전의 모델에서 급격하게 변하는 것을 막아야 한다고 합니다.

아래 실험 결과에서 볼 수 있듯이 ReST 를 사용했을 때 RLHF (Online RL) 를 사용하는 것 보다 더 나은 성능을 낼 수 있다고 합니다.

다만 iteration 을 너무 많이 반복하면 Reward hacking 이 발생한다는 리포트가 있습니다. 아래 그래프처럼 average reward 는 증가했지만 Initial model 과의 Human evaluation score 차이는 점점 줄어드는 것을 볼 수 있습니다. 따라서 적절한 iteration 을 실험적으로 찾아내는 작업이 반드시 필요합니다.

References

• Training language models to follow instructions with human feedback
• Learning to summarize from human feedback (Stiennon et al., 2020)
• RLAIF: Scaling Reinforcement Learning from Human Feedback with AI Feedback (Lee et al., 2023)
• Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena (Zheng et al., 2023)
• Reinforced Self-Training (ReST) for Language Modeling (Gulcehre et al., 2023)