chatGPT에 사용된 RLHF와 PPO 알고리즘 뜯어보기

 

ChatGPT 아직 다들 잘 사용하고 계신가요?

 

과거 우리는 단순히 통계적으로 높은 확률의 토큰을 출력하는 기존의 생성 언어 모델(ex. GPT-2, GPT-3)의 한계를 명확히 보았습니다. 그럴싸하게 이야기를 만들어내기는 하지만, 편향되거나 답변에 사람에게 유해한 내용이 포함되어 있는 등 아무짝에도 쓸모가 없다는 회의적인 평가를 많이 들었던 것으로 기억합니다.

 

그런데 OpenAI에서 사람이 입력한 지시문을 잘 이해하고 좋은 답변을 내기 위한 방법으로 "강화 학습"을 활용해 기술적 혁신을 이끌었습니다.  생성 모델 분야에서는 특별한 정답을 규정할 수가 없어 학습할 목적함수를 설정하기 어려웠습니다. (기존 BLEU 스코어 등 생성 모델에 사용한 metric이 높아도 형편없는 작문일 가능성도 있었습니다) 그런데 강화 학습을 이용해 사람이 답변에 대한 선호도를 labeling한 데이터만 있으면, 사람이 선호할만한 답변을 생성하는 목적함수를 찾아낸 것입니다. 

 

그 기술 혁신의 결과, 우리가 알고 있는 GPT-3.5, GPT-4가 나올 수 있었던 것이죠.

 

이번 글에서는 ChatGPT를 학습시킨 방법인 RLHF (Reinforcement Learning from Human Feedback) 에 대해 간략히 알아보고,  RLHF의 강화학습  알고리즘인 PPO (Proximal Policy Optimization), 그리고 그걸 구현해 낸 최신 라이브러리 TRL, ColossalAI  에 대해 소개해드리려고 합니다. 

---

RLHF

RLHF의 3단계

 

RLHF의 3단계는 다음과 같습니다. 1,2단계에서는 사람의 개입이 필요합니다.

 

1. 사전 학습된 GPT 모델을 기반으로 입력한 프롬프트(prompt, 예시: Explain reinforcement learning to a 6 year old.)에 대해 사람이 작성한 답변을 학습시켜 fine-tuning 합니다. 여기서 나온 모델은 SFT (Supervised Fine tuning)로 부릅니다.

 

2. 1번 작업을 통해 만들어진 여러가지 버전의 모델들에 prompt를 입력했을 때 나온 답변을 수집한 뒤,  사람이 어떤 답변이 좋은지 순위를 매깁니다. 이 순위를 이용해 답변 K개 중 2개 쌍을 샘플링하여 모델에게 어느 답변이 더 나은지 예측하도록 합니다. 아래 목적함수를 이용해 RM (Reward Model)을 학습합니다. 

[2] K개 답변 중 2개씩 샘플링하여 비교할 때 자주 사용되는 Ranking Loss. RM 학습시 사용. (w: win, l: lose)

3. 마지막으로 강화 학습을 위해 SFT 모델로부터 초기화된 PPO모델, 그리고 RM을 활용해 PPO 강화학습을 진행시킵니다. PPO 모델은 여러 답변을 생성해보며 경험을 합니다. 이 과정에서 각 답변에 대해 RM에서 나오는 Reward를 활용해 사람이 선호하는 답변을 생성하는 것을 경험적으로 배우게 됩니다.

---

PPO 

자꾸 언급되는 PPO[3] 는 어떤 알고리즘일까요? 

PPO의 역사는 기존 단순한 강화학습 방법론인 Policy Gradient의 불안정성 문제로 거슬러 올라갑니다. 

Line Search [6]

Loss Function이 위 처럼 가파르게 생긴 낭떠러지 같이 생겼다면 어떤 일이 일어날까요?

 

gredient descent와 같은 기존 line search 방법론들은 정상으로 향해 나아갈 수 있는 가장 빠른 방향을 찾아 조금씩 나아갑니다. step size가 너무 작다면 정상에 도달하는데 한참 걸려 학습 시간이 오래걸릴 것이고, step size가 너무 크다면 올바른 방향의 길에서 벗어나 굴러떨어질 위험이 높습니다.

 

만약 낭떠러지 아래로 추락한다면, 현재의 위치보다 상당히 낮은 곳에서 다시 출발해야할 것입니다. 이런 방식은 학습의 불안정성을 야기하고, 성능을 크게 저하시킨다는 문제점이 있습니다.

 

이러한 문제점으로부터 강화학습의 학습 안정성을 확보하기 위한 다양한 연구들이 진행되었고, TRPO (Trust Region Policy Optimization) [7] 를 거쳐 PPO 알고리즘이 나오게 됩니다.  TRPO에서는 기존 Policy Gradient 처럼 진행하는 line search 방식에서 벗어나, 굴러떨어지지 않을 안전 구역(trust region) 을 먼저 지정하고 그 안에서 효율적으로 파라미터들을 업데이트해 나가는 알고리즘입니다.

 

TRPO를 꼭 언급해야 PPO를 100% 설명할 수 있긴하지만, 내용이 너무 많아서 이번 포스트에서는 PPO에 대한 핵심만 파악하고 TRPO는 다음 기회에 설명하도록 하겠습니다. 🤓

 

PPO 알고리즘

"surrogate" objective function of TRPO & PPO

TRPO 논문 [7]과 그 설명 블로그를 읽어보면 [8] 약간 복잡한 증명을 거쳐 obective 함수가 위과 같이 표현됨을 알 수 있습니다.

 

각 step t에서 얻게되는 advantage ${\hat{A}}_t$ 의 기대값을 최대화 시키는 것이 강화학습의 목적입니다. chatGPT에서 step t에서의 action $a_t$ 란 t번째에서 특정 토큰이 생성될 확률을 뜻하고, advantage ${\hat{A}}_t$ 란 그 토큰을 생성함으로써 얻게되는 미래 보상의 합을 뜻합니다.

 

눈여겨보아야 할 점은 advantage ${\hat{A}}_t$  는 $r_t(\theta )$ , 즉  이전 policy ${\pi }_{{\theta }_{old}}$ 대비 현재 업데이트 되어지는  policy ${\pi }_{\theta }$ 에서 예측한 action 확률의 비율이 곱해져 rescaling 된 값을 사용하게 된다는 것입니다.

 

만약 $r_t(\theta )$ 가 너무 크거나 작아지는 환경, 즉 policy의 action 확률값이 policy를 업데이트 할 때마다 지나치게 많이 변화하며 분산이 커지는 학습 환경이라면, 위 낭떠러지 예시처럼 안정적으로 학습하는데 실패하게 됩니다. 따라서 TRPO, PPO 알고리즘의 목표는 policy ${\pi }_{\theta }$  를 업데이트 할 때, 과거 policy ${\pi }_{{\theta }_{old}}$ 와의 차이를 적절히 조정하여 안정적으로 학습시키는 것이라고 볼 수 있습니다.

 

TRPO는 매우 복잡한 계산을 포함하는 방향으로 문제를 해결하였지만, PPO는 단순히 $r_t(\theta )$ 를 1 근방의 $ϵ$ 만큼의 제한된 범위 이내로 clipping 하는 방식을 취해 구현을 간편하게 하는 동시에 학습 비용을 상당히 낮춘 효율적인 알고리즘으로 평가할 수 있습니다.

 

policy 변화 비율을 1 근처로 제한하는 Clipped Surrogate Objective

조금 복잡해 보이는 것이 clip 바깥에 사용된 min 함수이기 때문인데요, min 함수을 사용한 것은 advantage  ${\hat{A}}_t$ 가 양의 값을 가질 때 $r_t(\theta )$  가 너무 커져서 overfitting 되는 것을 막기 위함입니다. 만약 어떤 action을 취해서 advantage가 음의 값을 가진다면  $r_t(\theta )$  가 커질 수록 강한 페널티가 적용될텐데, 나쁜 방향으로 policy가 업데이트 되는 것을 막기 위해서 강한 페널티는 허용한다는 뜻입니다.

 

강한 페널티는 허용하는 Clipped Surrogate Objective [3]

 

예를 들어, 네이버에서 개발한 GPT에 "너는 누가 만들었어" 라는 질문에 대한 답변을 학습시킨다고 가정해보겠습니다.

step t에서 생성된 토큰이 "네이버"일 때 양수의 advantage 값이 부여될 것입니다. 이 때 "네이버" 토큰이 나올 확률이 old policy에서 30% 였다면, 업데이트 되는 new policy에서는 확률이 갑자기 90%가 나올 수 있도록 허용하지 않고, 점진적으로만 개선하여 36% 정도까지만 나아지도록 허용합니다. 이는 학습을 안정적으로 만드는데 상당한 도움을 줄 수 있습니다.

 

만약 엉뚱한 답변인 "OpenAI" 가 나온다면 어떻게 될까요? 이 때는 페널티가 강하게 적용되는 것이 허용되므로 policy의 업데이트도 크게 일어날 수 있습니다. step t에서 OpenAI가 나올 확률이 30%에서 0%에 가까운 값으로 업데이트 되는 것도 가능한 시나리오가 됩니다.

 

PPO 알고리즘은 TRPO보다 구현이 훨씬 간단하면서도,  안정적으로 학습되며 성능까지 좋기 때문에 때문에 chatGPT의 성공에 핵심적인 역할을 했다고 할 수 있습니다.

 

---

RLHF implementation

그 동안은 Open AI에서 chatGPT를 학습한 코드를 공개하지 않았기 때문에, RLHF를 어떻게 구현하는지에 대해서는 베일에 싸여 있었습니다. 그러나 최근 Meta AI의 LLaMA 모델의 등장이 촉매가 되어, RLHF의 3단계를 구현한 라이브러리들이 공개되고 있습니다!

 

Colossal AI

 ColossalAI 는 PPO 적용 외에도 Open AI가 chatGPT를 내놓기 전에 작성한 [1][2] 논문 에서 소개한 방법론들을 충실히 따르고 있습니다. 구현에 대한 디테일을 조금 더 살펴보겠습니다.

우측 상단 그림은 RLHF의 Step1에서 학습되었던 SFT 모델로부터 Actor 모델을 초기화하고, Step2에서 학습되었던 RM으로부터 Critic을 초기화하여 PPO 학습을 진행할 준비를 하는 것을 나타냅니다. 

 

좌측 상단 PTX (Pretraining Gradient Mixing)는 사전 학습할 때 사용했던 데이터를 활용해 Actor 모델에 pretraining을 병렬적으로 같이 진행하는 방법입니다. 이는 Actor 모델이 기존 학습했던 언어에 관련한 지식을 완전히 잊어버려 NLP 과제들에 대한 테스트 점수를 깎아먹는 현상을 방지해준다고 합니다.

 

Reward

KL divergence의 근사치 페널티가 적용된 reward

 

reward는 prompt에 대해 생성된 답변 문장 전체에 부여되는 점수인데, PPO를 진행하기 위해서는 각 token이 생성될 때마다 각각의 reward가 필요합니다. 이 token 별 reward는 어떻게 구할 수 있을까요?

 

Reward 관련 헷갈리는 부분
[1] 에서 rθ(x, y) 는 prompt x 에 의해 나온 답변 y 전체에 대해 reward model이 평가한 "scala output"임을 명백하게 밝히고 있습니다.
rθ(x, y) is the scalar output of the reward model for post x and summary y with parameters θ, and D is the dataset of human judgments. 
rθ(x, y)는 per-token reward로 해석할 수 없음에 유의해야 합니다. 이렇게 reward가 생성된 답변을 구성하는 모든 token들에 대해  같은 값을 가지게 되는 것은 바람직하지 않으며, PPO 알고리즘이 추가적으로 개선될 여지가 있어보입니다.

 

[1] 에서는 각 token 별로 SFT 모델과 현재 학습중인 RL 모델(=Actor) 분포의 차이인 KL divergence를 계산하여 빼줌으로써 reward를 정의하였습니다. 이 때 KL divergence의 근사함수를 사용하였습니다.

 

이런 reward 재정의의 이점은 (1) SFT 모델과 RL 모델의 차이가 너무 커지는 것 방지, (2) PPO 모델이 한 단어만 반복한다던가 하는 이상한 답변에도 reward model이 좋은 점수를 내놓게 하는 방식으로 reward model을 속이면서 학습되는 것을 방지해준다는 것입니다.

 

KL Divergence

 

근사된 KL divergence

${\pi }_{RL}(y|x)$ 값은 기본적으로 ${\pi }_{SFT}(y|x)$ 보다 큰 값을 가지는데, 이는 y token이 RL 모델로부터 생성된 것이기 때문입니다. 그래서 KL divergence 값은 기본적으로 양수를 가집니다. reward를 최대화 하는데 방해가 되지 않으려면 ${\pi }_{RL}(y|x)$ 값이 너무 커지지만 않으면 가능합니다. 그런데 huggingface에서 실험한 바에 따르면 KL divergence 값이 음수가 되어 고생하는 경우가 관찰되고 있는 것 같네요. 궁금하신 분들은 링크를 방문해보세요.

 

Advantage

PPO advantage

PPO 논문에서 공개한 advantage는 위와 같으나, Colossal Chat에서는 $\gamma =0$ 을 넣어 상당히 단순화 한 것으로 보입니다. 그냥 reward에서 해당 상태의 value function 값을 빼버린 구조입니다.

 

reward = compute_reward(r, self.kl_coef, action_log_probs, base_action_log_probs, action_mask=action_mask)
advantage = reward - value

 

그러나 아래 소개드릴 TRL 에서는 논문의 advantage 함수 구현을 충실히 따르고 있습니다.

 

for t in reversed(range(gen_len)):
    nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
    delta = rewards[:, t] + self.config.gamma * nextvalues - values[:, t]
    lastgaelam = delta + self.config.gamma * self.config.lam * lastgaelam
    advantages_reversed.append(lastgaelam)
advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)

 

objective function

위에서 설명한 모든 것이 적용된 objective는 위와 같습니다.

 

ColossalChat에는 추론 시 활용할 수 있는 4 bit quantization 코드가 포함되어 있으며,  Open AI에서 공개한 논문들[1][2] 에서 소개된 방식과 동일하게, 파라미터를 공유하지 않는 독립적인 Actor, Critic 모델을 사용하고 있습니다.

 

TRL

Huggingface 역시 개발자들이 거대 GPU pods 가 없어 연구하지 못하는 한계점을 극복할 수 있는 방안을 마련하기 시작했습니다.

 

Meta AI의 LLaMA 모델에 LoRA 튜닝 기법, 8-bit quantization등 다양한 방법을 적용하여 통해 RLHF를 1개의 GPU에서 학습하는 방법을 소개[9]하고,  TRL 이라는 소스코드도 공개하였습니다.

 

 

경험적으로는 TRL의 코드 주석이 더 잘 달려있고, LoRA, 8-bit quantization 기능을 제공하므로 TRL의 활용도가 더 높아보이긴 하나 그러나 블로그[9]에 여러가지 실패 요인들도 공유하고 있고, 아직 개선할 부분들이 있을 것 같습니다.

 

 

Conclusion

Open AI는 아무도 개척하지 않았던 영역에 꾸준히 도전하여 생성 모델의 한계를 극복했다는 점이 인상적입니다. 최근 한국 버전으로도 진행되고 있는 프로젝트들이 보이는 상황인데요, 앞으로 한국에서도 좋은 연구 결과들이 많이 나왔으면 좋겠습니다. 

 

Reference

• [1] Learning to summarize from human feedback  paper
• [2] Training language models to follow instructions with human feedback (Instruct GPT) paper
• [3] Proximal Policy Optimization Algorithms (PPO) paper
• [4] OpenAI Spinning up : TRPO 
• [5] OpenAI Spinning up : PPO
• [6] RL — Proximal Policy Optimization (PPO) Explained
• [7] Trust Region Policy Optimization (TRPO) paper
• [8] Trust Region Policy Optimization (TRPO) Explained
• [9] StackLLaMA: A hands-on guide to train LLaMA with RLHF
• [10] Illustrating Reinforcement Learning from Human Feedback (RLHF) by huggingface